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

base_model = "xlnet-base-cased"
session = load_variables()
sentences = load_sentences()
baseline_classifier = load_model(f"Dagobert42/{base_model}-biored-finetuned")
augmented_classifier = load_model(f"Dagobert42/{base_model}-biored-augmented")

st.title("Semantic Frame Augmentation")
st.subheader("Analysing challenging domains with only a handful of examples")

st.write(f"""This space uses different xlnet-base-cased models to find medical entities in a text.
This is a random sentence from [bigbio/biored](https://huggingface.co/datasets/bigbio/biored).
The sentence is tagged by a {base_model} model which has been trained on 200 exampled from the original dataset.
It should be making some mistakes.
""")

txt = sentences[session["counter"]]
if session["augment"]:
    st.write("Results on original biored data:")
    tokens = augmented_classifier(txt)
else:
    st.write("Results with data augmentation:")
    tokens = baseline_classifier(txt)

annotated_text(annotate_sentence(txt, tokens))

st.write("Now try the augmented model. Hopefully it's a bit better :)")
session["augment"] = st.toggle("augmentations on" if session["augment"] else "augmentations off")

st.write("Or load another sentence")
st.button( ":twisted_rightwards_arrows:", on_click=lambda: session.setdefault("counter", 0) + 1)