import torch
import streamlit as st
from transformers import pipeline
from random import choice

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

baseline_classifier = pipeline(
    model="Dagobert42/mobilebert-uncased-biored-finetuned-ner",
    task="ner",
    aggregation_strategy="simple"
    )
augmented_classifier = pipeline(
    model="Dagobert42/mobilebert-uncased-biored-augmented-ner",
    task="ner",
    aggregation_strategy="simple"
    )

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

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

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

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

st.subheader("Entity analysis:")
st.write(tokens)