import streamlit as st
from transformers import Pipeline, pipeline
import tokenizers
import torch
import json

@st.cache(
    hash_funcs={
        torch.nn.parameter.Parameter: lambda _: None,
        tokenizers.Tokenizer: lambda _: None,
        tokenizers.AddedToken: lambda _: None,
    },
    allow_output_mutation=True,
    show_spinner=False,
)

def load_engine() -> Pipeline:

    nlp_qa = pipeline(
        'question-answering',
        model='mrm8488/bert-italian-finedtuned-squadv1-it-alfa',
        tokenizer='mrm8488/bert-italian-finedtuned-squadv1-it-alfa'
    )

    return nlp_qa

with st.spinner(text="Sto preparando il necessario per rispondere alle tue domande personali..."):

    engine = load_engine()

with open("knowledge.json") as f:
    context = json.load(f)

st.title("La macchina che risponde alle tue domande personali")

input = st.text_input("Fai una domanda e comparirà la risposta!")

if input:
    try:

        answer = engine(
            {
                'question': input,
                'context': context["info"]
            }
        )

        st.subheader(answer["answer"])

    except:

        st.error("Qualcosa é andato storto. Prova di nuovo con un'altra domanda magari!")