Update app.py

app.py

@@ -1,46 +1,28 @@
 import gradio as gr
 from transformers import pipeline
 
-# Modelos pequenos e leves para rodar no ambiente gratuito
-model_a = pipeline("text-generation", model="gpt2")
-model_b = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")
+# Modelos geradores otimizados
+model_a = pipeline("text-generation", model="gpt2", return_full_text=False)
+model_b = pipeline("text-generation", model="EleutherAI/gpt-neo-125M", return_full_text=False)
 
-# Modelo juiz baseado em sentimento (positivo ou negativo)
+# Modelo juiz
 arbiter = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 
-# Função para julgar a melhor resposta com base em sentimento positivo
-def judge_response(prompt, response_a, response_b):
-    combined_a = f"Question: {prompt} Answer: {response_a}"
-    combined_b = f"Question: {prompt} Answer: {response_b}"
+# Julgamento por batch
+def judge_response(response_a, response_b):
+    results = arbiter([response_a, response_b])
+    score_a = results[0]['score']
+    score_b = results[1]['score']
+    return ("Model A", response_a) if score_a > score_b else ("Model B", response_b)
 
-    score_a = arbiter(combined_a)[0]['score']
-    score_b = arbiter(combined_b)[0]['score']
-
-    if score_a > score_b:
-        return "Model A", response_a
-    else:
-        return "Model B", response_b
-
-# Função principal do chatbot
+# Função principal
 def chatbot(prompt):
-    raw_response_a = model_a(prompt, max_new_tokens=60)[0]['generated_text']
-    raw_response_b = model_b(prompt, max_new_tokens=60)[0]['generated_text']
-
-    # Remove o prompt inicial da resposta
-    response_a = raw_response_a[len(prompt):].strip()
-    response_b = raw_response_b[len(prompt):].strip()
-
-    winner, final_response = judge_response(prompt, response_a, response_b)
+    response_a = model_a(prompt, max_new_tokens=40)[0]['generated_text'].strip()
+    response_b = model_b(prompt, max_new_tokens=40)[0]['generated_text'].strip()
+    winner, final_response = judge_response(response_a, response_b)
+    return prompt, response_a, response_b, winner, final_response
 
-    return (
-        prompt,
-        response_a,
-        response_b,
-        winner,
-        final_response
-    )
-
-# Interface Gradio
+# Interface
 iface = gr.Interface(
     fn=chatbot,
     inputs=gr.Textbox(label="Digite sua pergunta:"),
@@ -51,18 +33,8 @@ iface = gr.Interface(
         gr.Textbox(label="Modelo Vencedor"),
         gr.Textbox(label="Resposta Escolhida"),
     ],
-    title="Chatbot em Cascata com Julgamento",
-    description="""
-    Este chatbot utiliza dois modelos diferentes para responder à pergunta de um usuário.
-    Um terceiro modelo avalia qual resposta possui maior sentimento positivo e a apresenta como a melhor.
-
-    🧠 Modelos usados:
-    - Modelo A: GPT-2 (`gpt2`)
-    - Modelo B: GPT-Neo 125M (`EleutherAI/gpt-neo-125M`)
-    - Árbitro: DistilBERT SST-2 (`distilbert-base-uncased-finetuned-sst-2-english`)
-
-    Critério de julgamento: Resposta com maior sentimento positivo vence.
-    """
+    title="Chatbot com Julgamento",
+    description="Compara respostas de dois modelos e usa um árbitro para escolher a melhor com base no sentimento positivo."
 )
 
 iface.launch()