Refactor the app and change the context a bit

app.py

@@ -1,8 +1,12 @@
+import json
+
 import streamlit as st
-from transformers import Pipeline, pipeline
 import tokenizers
 import torch
-import json
+from transformers import Pipeline, pipeline
+
+from utils import get_answer
+
 
 @st.cache(
     hash_funcs={
@@ -13,18 +17,20 @@ import json
     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'
+        "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..."):
+
+with st.spinner(
+    text="Sto preparando il necessario per rispondere alle tue domande personali..."
+):
 
     engine = load_engine()
 
@@ -38,15 +44,11 @@ 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"])
+        answer = get_answer(input, context["info"], engine)
+        st.subheader(answer)
 
     except:
 
-        st.error("Qualcosa é andato storto. Prova di nuovo con un'altra domanda magari!")
+        st.error(
+            "Qualcosa é andato storto. Prova di nuovo con un'altra domanda magari!"
+        )

context.json

@@ -1,2 +1,2 @@
-{"info": "Sei Giuseppe, hai 69 anni e abiti assieme a Davide a Villafranca di Verona in Corso Vittorio Emanuele 238. Ora siamo a novembre dell'anno 2021 e hai avuto un'encefalite virale a gennaio 2021. Davide è tuo figlio e ha 37 anni. A causa dell'encefalite hai perso la memoria. Tua moglie Raffaella vive a Verona. Raffaella sta bene. Con Raffaella hai rapporti cordiali. I tuoi fratelli Alessandro, Giovanni e Grazia stanno bene."
+{"info": "Sei Giuseppe, hai 69 anni e abiti e vivi a Villafranca di Verona con Davide in Corso Vittorio Emanuele 238. Ora siamo a novembre dell'anno 2021 e hai avuto un'encefalite virale a gennaio 2021. Davide è tuo figlio, ha 37 anni. A causa dell'encefalite hai perso la memoria. Tua moglie Raffaella abita a Verona. Raffaella sta bene. Con Raffaella hai rapporti cordiali. Alessandro, Giovanni e Grazia sono i tuoi fratelli. I tuoi fratelli stanno bene."
 }

test_answers.py

@@ -0,0 +1,60 @@
+from utils import get_answer
+from transformers import pipeline
+import json
+
+nlp_qa = pipeline(
+    "question-answering",
+    model="mrm8488/bert-italian-finedtuned-squadv1-it-alfa",
+    tokenizer="mrm8488/bert-italian-finedtuned-squadv1-it-alfa",
+)
+
+with open("context.json") as f:
+    context = json.load(f)["info"]
+
+
+def test_name():
+    q = "Come mi chiamo?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "Giuseppe"
+
+
+def test_age():
+    q = "Quanti anni ho?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "69"
+
+
+def test_home():
+    q = "Dove vivo?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "Villafranca di Verona"
+
+
+def test_family():
+    q = "Con chi vivo?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "Davide"
+
+
+def test_family_2():
+    q = "Come si chiama mio figlio?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "Davide"
+
+
+def test_family_3():
+    q = "Quanti anni ha mio figlio?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "37"
+
+
+def test_family_4():
+    q = "Come sta Raffaella?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "Raffaella sta bene"
+
+
+def test_family_5():
+    q = "In che rapporti sono con Raffaella?"
+    a = get_answer(q, context, nlp_qa)
+    assert a == "cordiali"

utils.py

@@ -0,0 +1,5 @@
+def get_answer(input, context, engine):
+
+    answer = engine({"question": input, "context": context})
+
+    return answer["answer"]