Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForQuestionAnswering, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
 
 # Modelli
 MODEL_DEBERTA = "osiria/deberta-italian-question-answering"
@@ -10,10 +10,10 @@ tok_deb = AutoTokenizer.from_pretrained(MODEL_DEBERTA)
 mdl_deb = AutoModelForQuestionAnswering.from_pretrained(MODEL_DEBERTA)
 qa_deb = pipeline("question-answering", model=mdl_deb, tokenizer=tok_deb, device=-1)
 
-# Pipeline UmBERTo (generativo)
+# Pipeline UmBERTo (estrattivo)
 tok_umb = AutoTokenizer.from_pretrained(MODEL_UMBERTO)
-mdl_umb = AutoModelForCausalLM.from_pretrained(MODEL_UMBERTO)
-qa_umb = pipeline("text-generation", model=mdl_umb, tokenizer=tok_umb, device=-1)
+mdl_umb = AutoModelForQuestionAnswering.from_pretrained(MODEL_UMBERTO)
+qa_umb = pipeline("question-answering", model=mdl_umb, tokenizer=tok_umb, device=-1)
 
 def ensemble_invoice_qa(md_text: str, question: str):
     results = {}
@@ -28,22 +28,22 @@ def ensemble_invoice_qa(md_text: str, question: str):
     except Exception as e:
         results["DeBERTa (estrattivo)"] = {"errore": str(e)}
 
-    # Generativo (UmBERTo) con fix max_new_tokens
+    # Estrattivo (UmBERTo)
     try:
-        prompt = f"Domanda: {question}\nContesto: {md_text}\nRisposta:"
-        res_umb = qa_umb(prompt, max_new_tokens=64, do_sample=False)
-        results["UmBERTo (generativo)"] = {
-            "risposta": res_umb[0]["generated_text"].replace(prompt, "").strip()
+        res_umb = qa_umb(question=question, context=md_text)
+        results["UmBERTo (estrattivo)"] = {
+            "risposta": res_umb["answer"].strip(),
+            "confidenza": round(res_umb["score"], 3)
         }
     except Exception as e:
-        results["UmBERTo (generativo)"] = {"errore": str(e)}
+        results["UmBERTo (estrattivo)"] = {"errore": str(e)}
 
     return results
 
 # ================== UI Gradio ==================
 with gr.Blocks(theme=gr.themes.Base()) as demo:
     gr.Markdown("# 🧾 Invoice QA: Ensemble DeBERTa + UmBERTo")
-    gr.Markdown("Confronto tra risposte estrattive (DeBERTa) e generative (UmBERTo).")
+    gr.Markdown("Confronto tra risposte estrattive di due modelli italiani.")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -61,7 +61,7 @@ with gr.Blocks(theme=gr.themes.Base()) as demo:
             btn = gr.Button("🔍 Analizza Documento", variant="primary")
         
         with gr.Column(scale=1):
-            out_json = gr.JSON(label="Risultati Ensemble (Estrattivo vs Generativo)")
+            out_json = gr.JSON(label="Risultati Ensemble (DeBERTa vs UmBERTo)")
 
     btn.click(
         fn=ensemble_invoice_qa,