Update app.py

app.py

@@ -4,7 +4,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 # Modelos escolhidos
 model_name1 = "google/flan-t5-small"
 model_name2 = "google/flan-t5-base"
-arbitro_model_name = "google/flan-t5-small"  # Pode usar o mesmo ou outro
+arbitro_model_name = "google/flan-t5-base"  # Usando um modelo maior para arbitrar
 
 # Carregar tokenizadores e modelos
 tokenizer1 = AutoTokenizer.from_pretrained(model_name1)
@@ -17,21 +17,25 @@ tokenizer_arbitro = AutoTokenizer.from_pretrained(arbitro_model_name)
 model_arbitro = AutoModelForSeq2SeqLM.from_pretrained(arbitro_model_name)
 
 def gerar_resposta(model, tokenizer, pergunta):
-    # Prompt para focar na resposta correta da capital
-    prompt = f"Responda com o nome da capital do país nesta pergunta: {pergunta}"
-    inputs = tokenizer(prompt, return_tensors="pt")
+    # Formatar a pergunta para focar na capital
+    if "capital" not in pergunta.lower():
+        pergunta = f"Qual é a capital do {pergunta}?"
+    
+    inputs = tokenizer(pergunta, return_tensors="pt")
     outputs = model.generate(**inputs, max_length=20)
     resposta = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return resposta.strip()
 
 def arbitro(pergunta, resp1, resp2):
-    # Prompt para o árbitro escolher a melhor resposta
+    # Melhorar o prompt para o árbitro
     prompt = (
+        f"Decida qual resposta está mais correta para esta pergunta sobre geografia:\n"
         f"Pergunta: {pergunta}\n"
         f"Resposta 1: {resp1}\n"
-        f"Resposta 2: {resp2}\n"
-        "Qual resposta é mais correta? Responda apenas 1 ou 2."
+        f"Resposta 2: {resp2}\n\n"
+        f"Responda apenas com o número 1 ou 2, sem explicações."
     )
+    
     inputs = tokenizer_arbitro(prompt, return_tensors="pt")
     outputs = model_arbitro.generate(**inputs, max_length=5)
     escolha = tokenizer_arbitro.decode(outputs[0], skip_special_tokens=True).strip()
@@ -41,29 +45,35 @@ def arbitro(pergunta, resp1, resp2):
     elif escolha == "2":
         return resp2, "Modelo 2 (flan-t5-base)"
     else:
-        # Caso o árbitro não dê 1 ou 2, escolhe a resposta 1 por padrão
-        return resp1, "Modelo 1 (flan-t5-small)"
+        # Se não conseguir decidir, escolhe a mais longa (heurística simples)
+        return (resp1, "Modelo 1 (flan-t5-small)") if len(resp1) > len(resp2) else (resp2, "Modelo 2 (flan-t5-base)")
 
 def chatbot(pergunta):
+    # Garantir que a pergunta está formatada corretamente
+    pergunta = pergunta.strip().lower()
+    if not pergunta.endswith("?"):
+        pergunta += "?"
+    
     resposta1 = gerar_resposta(model1, tokenizer1, pergunta)
     resposta2 = gerar_resposta(model2, tokenizer2, pergunta)
     resposta_final, modelo_escolhido = arbitro(pergunta, resposta1, resposta2)
 
     return (
-        f"Resposta selecionada:\n{resposta_final}\n\nModelo escolhido:\n{modelo_escolhido}",
-        f"Resposta Modelo 1 (flan-t5-small):\n{resposta1}",
-        f"Resposta Modelo 2 (flan-t5-base):\n{resposta2}"
+        f"Resposta selecionada: {resposta_final}\nModelo escolhido: {modelo_escolhido}",
+        f"Modelo 1 (flan-t5-small): {resposta1}",
+        f"Modelo 2 (flan-t5-base): {resposta2}"
     )
 
 iface = gr.Interface(
     fn=chatbot,
-    inputs=gr.Textbox(label="Digite uma pergunta sobre capitais"),
+    inputs=gr.Textbox(label="Digite um país ou pergunta sobre capitais", placeholder="Ex: Brasil ou Qual a capital do Brasil?"),
     outputs=[
-        gr.Textbox(label="Resposta selecionada"),
-        gr.Textbox(label="Resposta Modelo 1"),
-        gr.Textbox(label="Resposta Modelo 2")
+        gr.Textbox(label="Resposta Final"),
+        gr.Textbox(label="Resposta do Modelo Pequeno"),
+        gr.Textbox(label="Resposta do Modelo Base")
     ],
-    title="Chatbot em Cascata - Perguntas sobre Capitais"
+    title="Chatbot em Cascata - Perguntas sobre Capitais",
+    description="Digite um país ou pergunta sobre capitais para comparar dois modelos de IA"
 )
 
 if __name__ == "__main__":