import streamlit as st
from annotated_text import annotated_text
from transformers import pipeline
from PIL import Image
import re

st.sidebar.header("**Instructions**")
st.sidebar.markdown("Démonstrateur du modèle [Camembert-NER](CATIE-AQ/Camembert-NER-base-frenchNER) entraîné sur plus de 425 000 données en français. Le modèle est capable d'étiquetter les entités LOC (Localisation), PER (Personne), ORG (Organisation) dans le texte fourni. Il est disponible en version *base* (110M de paramètres) et *large* (335M de paramètres). Pour l'essayer, sélectionnez la version de votre choix ci-dessous, puis renseignez un texte. Enfin appuyez sur le bouton « Appliquer le modèle » pour observer la réponse trouvée par le modèle. Pour en savoir plus sur ce modèle, vous pouvez lire l'[article de blog](https://blog.vaniila.ai/NER/) détaillant la démarche suvie.")
version = st.sidebar.radio("Choix de la version du modèle :", ["version base", "version large"])
st.sidebar.markdown("---")
st.sidebar.markdown("Ce modèle a été entraîné via la plateforme [*Vaniila*](https://www.vaniila.ai/) du [CATIE](https://www.catie.fr/).")

image_path = 'Vaniila.png'
image = Image.open(image_path)
st.sidebar.image(image, caption=None, width=None, use_column_width=None, clamp=False, channels="RGB", output_format="auto")

@st.cache_resource
def load_model(version,text):
    if version == "version base":
        ner = pipeline('token-classification', model='bourdoiscatie/camembert-base-frenchNER_3entities', tokenizer='bourdoiscatie/camembert-base-frenchNER_3entities', grouped_entities=True)
        result = ner(text)
        return result   
    else:
        ner = pipeline('token-classification', model='bourdoiscatie/camembert-base-frenchNER_4entities', tokenizer='bourdoiscatie/camembert-base-frenchNER_4entities', aggregation_strategy="max")
        result = ner(text)
        return result   

def getcolor(texts, labels):
    colors = {'LOC': '#38419D', 'PER': '#BF3131', 'ORG': '#597E52', 'MISC':'#F1C232'}
    return [(t,l,colors[l]) for t, l in zip(texts, labels)]

def color_annotation(to_print,transcript) :  
    text_ner = []
    label_ner = []
    for i in range(len(to_print)) :
        text_ner.append(to_print[i]["word"])
        label_ner.append(to_print[i]["entity_group"])

    anns = getcolor(text_ner, label_ner)
    anns = list(set(anns))
    text_ner = list(set(text_ner))
    text_ner = list(sorted(text_ner, key = len))

    display = transcript.replace(" ",' ","')
    for i in range(len(anns)):
        for j in range(len(text_ner)):
            if text_ner[j] == anns[i][0]:
                 display = display.replace(text_ner[j].replace(" ",' ","'),str(anns[i]))
    display = display.replace(', ","',',').replace(".","")
    
    for i in re.findall(r"\((.*?)\)", display) : # pour gérer les cas de mots inclus dans des n_grams
        if "(" in i:
            display = display.replace(i+")",i.split(",")[0].replace("('","")).replace("''","'")
            
    # sanity check
    display = display.replace(")",')","').replace(')","","',')","').replace("(",'","(').replace('","","(','","(')

    return display




            
st.markdown("<h2 style='text-align: center'>Camembert-NER", unsafe_allow_html=True)    
st.markdown("<h4 style='text-align: center'>"+version, unsafe_allow_html=True)    
option = st.selectbox(
    'Choix du mode',
    ('Texte libre', 'Exemple'))

if option == "Exemple":
    text = st.text_area("Votre texte", value="Assurés de disputer l'Euro 2024 en Allemagne l'été prochain (du 14 juin au 14 juillet) depuis leur victoire aux Pays-Bas, les Bleus ont fait le nécessaire pour avoir des certitudes. Avec six victoires en six matchs officiels et un seul but encaissé, Didier Deschamps a consolidé les acquis de la dernière Coupe du monde. Les joueurs clés sont connus : Kylian Mbappé, Aurélien Tchouameni, Antoine Griezmann, Ibrahima Konaté ou encore Mike Maignan.",height=175)
    if text:
        to_print = load_model(version,text)
        display = color_annotation(to_print,text)
        list_to_display = [] # pour pouvoir afficher la couleur, on doit passer les mots à colorier de str en tuple
        for i in range(len(display.split('","'))):
            if "#" in display.split('","')[i]:
                list_to_display.append(eval(display.split('","')[i]))
            else :
                list_to_display.append(display.split('","')[i])
        annotated_text(*list_to_display)
        st.write("\n")
        with st.expander("Afficher le score pour chacune des entitées trouvées :"):
            for i in range(len(to_print)) :
                st.write("- Score pour que ",to_print[i]["word"]," soit de type", to_print[i]["entity_group"]," : ",round(to_print[i]["score"],3))
    
else:
    text = st.text_area("Votre texte", value="",height=175)
    if text:
        col1, col2, col3 = st.columns(3)
        if col2.button('Appliquer le modèle'):
            to_print = load_model(version,text)
            display = color_annotation(to_print,text)
            list_to_display = [] # pour pouvoir afficher la couleur, on doit passer les mots à colorier de str en tuple
            for i in range(len(display.split('","'))):
                if "#" in display.split('","')[i]:
                    list_to_display.append(eval(display.split('","')[i]))
                else :
                    list_to_display.append(display.split('","')[i])
            annotated_text(*list_to_display)
            st.write("\n")
            with st.expander("Afficher le score pour chacune des entitées trouvées :"):
                for i in range(len(to_print)) :
                    st.write("- Score pour que ",to_print[i]["word"]," soit de type", to_print[i]["entity_group"]," : ",round(to_print[i]["score"],3))