from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
import sentencepiece
import streamlit as st
import pandas as pd

text_1 = "ddd"

text_2 = """ddd"""

st.title("Demo for Biomedical POS Tagging in French with DrBERT")
st.sidebar.write("Model : DrBERT-7GB base CAS corpus POS tagging")
st.sidebar.write("For details of model: 'https://huggingface.co/Dr-BERT/DrBERT-7GB'")

model_checkpoint = "Dr-BERT/DrBERT-7GB"
aggregation = "simple"

st.subheader("Select Text")
context_1 = st.text_area("Text #1", text_1, height=128)
context_2 = st.text_area("Text #2", text_2, height=128)
context_3 = st.text_area("New Text", value="", height=128)

context = st.radio("Select Text", ("Text #1", "Text #2", "New Text"))

if context == "Text #1":
    input_text = context_1
elif context == "Text #2":
    input_text = context_2
elif context == "New Text":
    input_text = context_3

@st.cache(allow_output_mutation=True)
def setModel(model_checkpoint, aggregation):
    model = AutoModelForTokenClassification.from_pretrained(model_checkpoint)
    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
    return pipeline('token-classification', model=model, tokenizer=tokenizer, aggregation_strategy=aggregation)
        
Run_Button = st.button("Run", key=None)
if Run_Button == True:
    
    ner_pipeline = setModel(model_checkpoint, aggregation)
    output = ner_pipeline(input_text)

    df = pd.DataFrame.from_dict(output)

    if aggregation != "none":
        df.rename(index=str,columns={'entity_group':'POS Tag'},inplace=True)
    else:
        df.rename(index=str,columns={'entity_group':'POS Tag'},inplace=True)
    
    cols_to_keep = ['word','POS Tag','score','start','end']
    df_final = df[cols_to_keep]
    
    st.subheader("POS Tags")
    st.dataframe(df_final)