import torch
import streamlit as st
from random import choice
from annotated_text import annotated_text
from helpers import *

with open("sentences.pt", 'rb') as f:
    sentences = torch.load(f)
sentence = choice(sentences)

st.title("Semantic Frame Augmentation")
st.subheader("Analysing difficult low-resource domains with only a handful of examples")

st.write("This space uses a google/mobilebert-uncased model for NER")
augment = st.toggle('Use augmented model for NER', value=False)

txt = st.text_area(
    "Text to analyze",
    sentence,
    max_chars=500
    )

if augment:
    st.write("with augmentation:")
    tokens = augmented_classifier(txt)
else:
    st.write("without augmentation:")
    tokens = baseline_classifier(txt)

st.subheader("Entity analysis:")
annotated_text(annotate_sentence(sentence, tokens))