Update app.py

app.py

@@ -1,26 +1,26 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import MT5Tokenizer, MT5ForConditionalGeneration
+import gradio as gr
 
-# Modelos A, B e Árbitro (modelo C)
-model_a_name = "google/flan-t5-base"
+# Modelos
+model_a_name = "google/flan-t5-small"
 model_b_name = "google/mt5-small"
-judge_model_name = "google/flan-t5-base"
+model_j_name = "google/flan-t5-small"  # Árbitro
 
-# Carrega modelos e tokenizers
+# Tokenizers
 tokenizer_a = AutoTokenizer.from_pretrained(model_a_name)
-model_a = AutoModelForSeq2SeqLM.from_pretrained(model_a_name)
-
-from transformers import MT5Tokenizer
 tokenizer_b = MT5Tokenizer.from_pretrained(model_b_name, use_fast=False)
-model_b = AutoModelForSeq2SeqLM.from_pretrained(model_b_name)
+tokenizer_j = AutoTokenizer.from_pretrained(model_j_name)
 
-tokenizer_j = AutoTokenizer.from_pretrained(judge_model_name)
-model_j = AutoModelForSeq2SeqLM.from_pretrained(judge_model_name)
+# Modelos carregados
+model_a = AutoModelForSeq2SeqLM.from_pretrained(model_a_name)
+model_b = MT5ForConditionalGeneration.from_pretrained(model_b_name)
+model_j = AutoModelForSeq2SeqLM.from_pretrained(model_j_name)
 
 def gerar_resposta(model, tokenizer, prompt):
     prompt_instruido = f"Question: {prompt}\nAnswer:"
-    inputs = tokenizer(prompt_instruido, return_tensors="pt")
+    inputs = tokenizer(prompt_instruido, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -33,12 +33,9 @@ def gerar_resposta(model, tokenizer, prompt):
     return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
 def limpar_resposta(resposta):
-    if "<extra_id" in resposta.lower():
+    if "<extra_id" in resposta.lower() or resposta.strip() == "":
         return ""
-    return resposta
-
-resp_a = limpar_resposta(gerar_resposta(model_a, tokenizer_a, prompt))
-resp_b = limpar_resposta(gerar_resposta(model_b, tokenizer_b, prompt))
+    return resposta.strip()
 
 def julgar_respostas(prompt, resp_a, resp_b):
     prompt_julgamento = (
@@ -52,26 +49,38 @@ def julgar_respostas(prompt, resp_a, resp_b):
         return "B"
     return "A"
 
-def processar(prompt):
+def responder(prompt):
+    # Geração
     resp_a = gerar_resposta(model_a, tokenizer_a, prompt)
     resp_b = gerar_resposta(model_b, tokenizer_b, prompt)
+
+    # Limpeza
+    resp_a = limpar_resposta(resp_a)
+    resp_b = limpar_resposta(resp_b)
+
+    # Se ambas estão vazias
+    if not resp_a and not resp_b:
+        return "⚠️ Nenhuma resposta válida foi gerada."
+
     melhor = julgar_respostas(prompt, resp_a, resp_b)
+    resposta_final = resp_a if melhor == "A" else resp_b
 
-    final = resp_a if melhor == "A" else resp_b
-    saida = (
-        f"🟡 Prompt: {prompt}\n\n"
-        f"🔹 Resposta A (Flan-T5):\n{resp_a}\n\n"
-        f"🔸 Resposta B (mT5):\n{resp_b}\n\n"
-        f"✅ Melhor resposta: Modelo {melhor} selecionado pelo árbitro.\n\n"
-        f"💬 Resposta final:\n{final}"
-    )
-    return saida
+    return f"""🟡 Prompt: {prompt}
+
+🔹 Resposta A (Flan-T5):
+{resp_a or '[Resposta inválida]'}
+
+🔸 Resposta B (mT5):
+{resp_b or '[Resposta inválida]'}
+
+✅ Melhor resposta: Modelo {melhor} selecionado pelo árbitro.
+
+💬 Resposta final:
+{resposta_final or '[Nenhuma resposta válida]'}
+"""
 
 # Interface Gradio
-gr.Interface(
-    fn=processar,
-    inputs=gr.Textbox(lines=2, placeholder="Digite sua pergunta...", label="Prompt"),
-    outputs=gr.Textbox(label="Resposta Final"),
-    title="Chatbot em Cascata com Árbitro",
-    description="Dois modelos geram respostas e um terceiro escolhe a melhor. Rodando localmente sem API externa."
-).launch()
+iface = gr.Interface(fn=responder, inputs="text", outputs="text", title="Chatbot em Cascata com Árbitro")
+
+if __name__ == "__main__":
+    iface.launch()