Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 # Modelos
 model_name1 = "google/flan-t5-small"
 model_name2 = "google/flan-t5-base"
-arbitro_model_name = "google/flan-t5-base"
+arbitro_model_name = "google/flan-t5-large"  # Modelo maior para arbitragem
 
 # Carregar modelos e tokenizadores
 tokenizer1 = AutoTokenizer.from_pretrained(model_name1)
@@ -17,11 +17,29 @@ model2 = AutoModelForSeq2SeqLM.from_pretrained(model_name2)
 tokenizer_arbitro = AutoTokenizer.from_pretrained(arbitro_model_name)
 model_arbitro = AutoModelForSeq2SeqLM.from_pretrained(arbitro_model_name)
 
+# Lista de capitais conhecidas para validação
+CAPITAIS_CONHECIDAS = {
+    'brasil': 'Brasília',
+    'alemanha': 'Berlim',
+    'frança': 'Paris',
+    'japão': 'Tóquio',
+    'itália': 'Roma',
+    'espanha': 'Madri',
+    'portugal': 'Lisboa',
+    'argentina': 'Buenos Aires',
+    'estados unidos': 'Washington',
+    'canadá': 'Ottawa',
+    # Adicione mais conforme necessário
+}
+
 # Geração da resposta de cada modelo com prompt reforçado
 def gerar_resposta(model, tokenizer, pergunta):
     prompt = (
-        "Answer ONLY with the name of the capital city for the following question.\n"
-        "Do NOT answer with the country name.\n"
+        "I will ask you about capital cities. Always respond with just the capital name.\n"
+        "Examples:\n"
+        "Question: What is the capital of France? Answer: Paris\n"
+        "Question: Capital of Japan? Answer: Tokyo\n"
+        "Question: Qual é a capital do Brasil? Answer: Brasília\n"
         f"Question: {pergunta}\n"
         "Answer:"
     )
@@ -32,26 +50,56 @@ def gerar_resposta(model, tokenizer, pergunta):
 
 # 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()))
+    # Verifica se está na lista de capitais conhecidas
+    resposta_limpa = resposta.strip().lower()
+    for capital in CAPITAIS_CONHECIDAS.values():
+        if capital.lower() == resposta_limpa:
+            return True
+    return False
+
+# Correção de respostas óbvias
+def corrigir_resposta(pergunta, resposta):
+    pergunta_lower = pergunta.lower()
+    resposta_lower = resposta.lower()
+    
+    # Verifica se a resposta é igual ao país mencionado na pergunta
+    for pais, capital in CAPITAIS_CONHECIDAS.items():
+        if pais in pergunta_lower and resposta_lower == pais:
+            return capital
+    
+    # Correções específicas
+    if "brasil" in pergunta_lower and resposta_lower == "brasil":
+        return "Brasília"
+    
+    return resposta
 
 # Árbitro decide qual resposta é melhor
 def arbitro(pergunta, resp1, resp2):
+    # Primeiro verifica se alguma resposta está na lista de capitais conhecidas
+    resp1_corrigida = corrigir_resposta(pergunta, resp1)
+    resp2_corrigida = corrigir_resposta(pergunta, resp2)
+    
+    if eh_capital_valida(resp1_corrigida) and not eh_capital_valida(resp2_corrigida):
+        return resp1_corrigida, "Modelo 1 (corrigido)"
+    elif eh_capital_valida(resp2_corrigida) and not eh_capital_valida(resp1_corrigida):
+        return resp2_corrigida, "Modelo 2 (corrigido)"
+    
+    # Se ambas ou nenhuma for válida, usa o árbitro
     prompt = (
-        "You are a geography expert.\n"
+        "You are a geography expert. Choose the correct capital city.\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."
+        f"Option 1: {resp1_corrigida}\n"
+        f"Option 2: {resp2_corrigida}\n"
+        "Which option 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 == "2":
-        return resp2, "Modelo 2 (flan-t5-base)"
+        return resp2_corrigida, "Modelo 2 (flan-t5-base)"
     else:
-        return resp1, "Modelo 1 (flan-t5-small)"
+        return resp1_corrigida, "Modelo 1 (flan-t5-small)"
 
 # Função principal do chatbot
 def chatbot(pergunta):
@@ -59,7 +107,7 @@ def chatbot(pergunta):
     resposta2 = gerar_resposta(model2, tokenizer2, pergunta)
     resposta_final, modelo_escolhido = arbitro(pergunta, resposta1, resposta2)
 
-    # Validação
+    # Validação final
     if not eh_capital_valida(resposta_final):
         resposta_final = "Não consegui identificar a capital corretamente."
 
@@ -78,7 +126,7 @@ iface = gr.Interface(
         gr.Textbox(label="Resposta Modelo 1"),
         gr.Textbox(label="Resposta Modelo 2")
     ],
-    title="Chatbot em Cascata - Perguntas sobre Capitais (sem listas)",
+    title="Chatbot em Cascata - Perguntas sobre Capitais (Melhorado)",
     description="Insira uma pergunta como 'Qual é a capital da Alemanha?' e veja como os modelos escolhem a melhor resposta."
 )