Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
+import re
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Modelos escolhidos
+# Modelos
 model_name1 = "google/flan-t5-small"
 model_name2 = "google/flan-t5-base"
-arbitro_model_name = "google/flan-t5-base"  # Usando um modelo maior para arbitrar
+arbitro_model_name = "google/flan-t5-base"
 
-# Carregar tokenizadores e modelos
+# Carregar modelos e tokenizadores
 tokenizer1 = AutoTokenizer.from_pretrained(model_name1)
 model1 = AutoModelForSeq2SeqLM.from_pretrained(model_name1)
 
@@ -16,65 +17,70 @@ model2 = AutoModelForSeq2SeqLM.from_pretrained(model_name2)
 tokenizer_arbitro = AutoTokenizer.from_pretrained(arbitro_model_name)
 model_arbitro = AutoModelForSeq2SeqLM.from_pretrained(arbitro_model_name)
 
+# Geração da resposta de cada modelo com prompt reforçado
 def gerar_resposta(model, tokenizer, pergunta):
-    # 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")
+    prompt = (
+        "Answer ONLY with the name of the capital city for the following question.\n"
+        "Do NOT answer with the country name.\n"
+        f"Question: {pergunta}\n"
+        "Answer:"
+    )
+    inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(**inputs, max_length=20)
     resposta = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return resposta.strip()
 
+# Função para validar se parece uma capital
+def eh_capital_valida(resposta):
+    # Considera resposta válida se tem 1 ou 2 palavras, só letras e espaços
+    return bool(re.match(r"^[A-Za-zÀ-ÿ\s]{2,40}$", resposta.strip()))
+
+# Árbitro decide qual resposta é melhor
 def arbitro(pergunta, resp1, resp2):
-    # 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\n"
-        f"Responda apenas com o número 1 ou 2, sem explicações."
+        "You are a geography expert.\n"
+        f"Question: {pergunta}\n"
+        f"Answer 1: {resp1}\n"
+        f"Answer 2: {resp2}\n"
+        "Which answer is the correct capital? Reply only with 1 or 2."
     )
-    
     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()
 
-    if escolha == "1":
-        return resp1, "Modelo 1 (flan-t5-small)"
-    elif escolha == "2":
+    if escolha == "2":
         return resp2, "Modelo 2 (flan-t5-base)"
     else:
-        # 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)")
+        return resp1, "Modelo 1 (flan-t5-small)"
 
+# Função principal do chatbot
 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)
 
+    # Validação
+    if not eh_capital_valida(resposta_final):
+        resposta_final = "Não consegui identificar a capital corretamente."
+
     return (
-        f"Resposta selecionada: {resposta_final}\nModelo escolhido: {modelo_escolhido}",
-        f"Modelo 1 (flan-t5-small): {resposta1}",
-        f"Modelo 2 (flan-t5-base): {resposta2}"
+        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}"
     )
 
+# Interface Gradio
 iface = gr.Interface(
     fn=chatbot,
-    inputs=gr.Textbox(label="Digite um país ou pergunta sobre capitais", placeholder="Ex: Brasil ou Qual a capital do Brasil?"),
+    inputs=gr.Textbox(label="Digite uma pergunta sobre capitais"),
     outputs=[
-        gr.Textbox(label="Resposta Final"),
-        gr.Textbox(label="Resposta do Modelo Pequeno"),
-        gr.Textbox(label="Resposta do Modelo Base")
+        gr.Textbox(label="Resposta selecionada"),
+        gr.Textbox(label="Resposta Modelo 1"),
+        gr.Textbox(label="Resposta Modelo 2")
     ],
-    title="Chatbot em Cascata - Perguntas sobre Capitais",
-    description="Digite um país ou pergunta sobre capitais para comparar dois modelos de IA"
+    title="Chatbot em Cascata - Perguntas sobre Capitais (sem listas)",
+    description="Insira uma pergunta como 'Qual é a capital da Alemanha?' e veja como os modelos escolhem a melhor resposta."
 )
 
-if __name__ == "__main__":
+if _name_ == "_main_":
     iface.launch()