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

st.sidebar.header("**Instructions**")
st.sidebar.markdown("Démonstrateur du modèle [QAmembert](https://huggingface.co/collections/CATIE-AQ/catie-french-qa-pack-650821750f44c341cdb8ec91) entraîné sur plus de 220 000 contextes/questions/réponses en français et capable d'indiquer si la réponse se situe ou non dans le texte fourni. Il est disponible en quatre versions : QAmemBERT1-base (110M de paramètres et taille de contexte de 512 tokens), QAmemBERT2 (112M, 1024 tokens), QAmemBERTa (111M, 1024 tokens) (**ce modèle donne les meilleurs résultats d'après nos benchmarks**) et QAmemBERT1-large (336M, 512 tokens).  Pour utiliser l'application, sélectionnez la version de votre choix ci-dessous, puis renseignez un texte et une question à propos de celui-ci. 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/QA/) détaillant la démarche suvie.")
version = st.sidebar.radio("Choix de la version du modèle :", ["QAmemBERT1-base", "QAmemBERT2", "QAmemBERTa", "QAmemBERT1-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):
    if version == "QAmemBERT1-base":
        return pipeline('question-answering', model='CATIE-AQ/QAmembert', tokenizer='CATIE-AQ/QAmembert')
    if version == "QAmemBERT2":
        return pipeline('question-answering', model='CATIE-AQ/QAmembert2', tokenizer='CATIE-AQ/QAmembert2')
    if version == "QAmemBERTa":
        return pipeline('question-answering', model='CATIE-AQ/QAmemberta', tokenizer='CATIE-AQ/QAmemberta')
    else:
        return pipeline('question-answering', model='CATIE-AQ/QAmembert-large', tokenizer='CATIE-AQ/QAmembert-large')
        
def results():
    col1, col2, col3 = st.columns(3)
    if col2.button('Appliquer le modèle'):
        result = model(question,context)
        st.write("Résultat")
        if result['score'] < 0.01:
            st.info("La réponse n'est pas dans le contexte fourni.")
        else :
            st.info(result['answer'])
        with st.expander("Voir le détail :"):
            st.write("• Score :", result['score'])
            st.write("• Début de la réponse :", result['start'])
            st.write("• Fin de la réponse :", result['end'])
            st.write("• Réponse :", result['answer'])
    pass
            
            
st.markdown("<h2 style='text-align: center'>QAmembert", unsafe_allow_html=True)    
st.markdown("<h4 style='text-align: center'>"+version, unsafe_allow_html=True)    
model = load_model(version)

option = st.selectbox(
    'Choix du mode',
    ('Texte libre', 'Exemple 1', 'Exemple 2'))

if option == "Exemple 1":
    context = st.text_area("Votre texte", value="Le français est une langue indo-européenne de la famille des langues romanes dont les locuteurs sont appelés francophones. Elle est parfois surnommée la langue de Molière.  Le français est parlé, en 2023, sur tous les continents par environ 321 millions de personnes : 235 millions l'emploient quotidiennement et 90 millions en sont des locuteurs natifs. En 2018, 80 millions d'élèves et étudiants s'instruisent en français dans le monde. Selon l'Organisation internationale de la francophonie (OIF), il pourrait y avoir 700 millions de francophones sur Terre en 2050.",height=175)
    question = st.text_input("Votre question", value="Combien de personnes utilisent le français tous les jours ?")
    if context and question:
        results()

elif option == "Exemple 2":
    context = st.text_area("Votre texte", value="Le français est une langue indo-européenne de la famille des langues romanes dont les locuteurs sont appelés francophones. Elle est parfois surnommée la langue de Molière.  Le français est parlé, en 2023, sur tous les continents par environ 321 millions de personnes : 235 millions l'emploient quotidiennement et 90 millions en sont des locuteurs natifs. En 2018, 80 millions d'élèves et étudiants s'instruisent en français dans le monde. Selon l'Organisation internationale de la francophonie (OIF), il pourrait y avoir 700 millions de francophones sur Terre en 2050.",height=175)
    question = st.text_input("Votre question", value="Quel est le meilleur vin du monde ?")
    if context and question:
        results()
    
else:
    context = st.text_area("Votre texte", value="",height=175)
    question = st.text_input("Votre question", value="")
    if context and question:
        results()