Update app.py

app.py

@@ -1,146 +1,68 @@
 import gradio as gr
-import torch
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 
-# Configurações de memória
-DEVICE = 0 if torch.cuda.is_available() else -1
-TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
+# 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
-MODELS = {
-    "GPT-2": {
-        "name": "gpt2",
-        "max_tokens": 150
-    },
-    "GPT-Neo 125M": {
-        "name": "EleutherAI/gpt-neo-125M",
-        "max_tokens": 150
-    }
-}
+# Modelo juiz baseado em sentimento (positivo ou negativo)
+arbiter = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 
-# Carrega os modelos de geração
-loaded_models = {}
-for model_name, config in MODELS.items():
-    try:
-        model = AutoModelForCausalLM.from_pretrained(
-            config["name"],
-            torch_dtype=TORCH_DTYPE,
-            device_map="auto" if DEVICE == 0 else None,
-            low_cpu_mem_usage=True
-        )
-        tokenizer = AutoTokenizer.from_pretrained(config["name"])
-        pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            device=DEVICE,
-            return_full_text=False,
-            pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id else 50256
-        )
-        loaded_models[model_name] = {
-            "pipe": pipe,
-            "max_tokens": config["max_tokens"]
-        }
-        print(f"✅ {model_name} carregado com sucesso")
-    except Exception as e:
-        print(f"❌ Erro ao carregar {model_name}: {str(e)}")
-        loaded_models[model_name] = None
+# 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}"
 
-# Modelo árbitro para classificar as respostas
-try:
-    arbiter = pipeline(
-        "text-classification",
-        model="distilbert-base-uncased-finetuned-sst-2-english",
-        device=DEVICE
-    )
-    print("✅ Modelo árbitro carregado com sucesso")
-except Exception as e:
-    print(f"❌ Erro ao carregar modelo árbitro: {str(e)}")
-    arbiter = None
-
-def format_prompt(user_input):
-    return f"Responda de forma clara e concisa: {user_input.strip()}"
+    score_a = arbiter(combined_a)[0]['score']
+    score_b = arbiter(combined_b)[0]['score']
 
-def generate_response(prompt, model_name):
-    if model_name not in loaded_models or not loaded_models[model_name]:
-        return "Modelo não disponível"
-    config = loaded_models[model_name]
-    try:
-        response = config["pipe"](
-            format_prompt(prompt),
-            max_new_tokens=config["max_tokens"],
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.2
-        )[0]['generated_text'].strip()
-        return response if response else "Nenhuma resposta gerada."
-    except Exception as e:
-        return f"Erro na geração: {str(e)}"
-
-def judge_responses(resp1, resp2):
-    if not arbiter:
-        # Se não houver árbitro, retorne a primeira resposta por padrão
-        return ("GPT-2", resp1)
-    inputs = [resp1, resp2]
-    try:
-        results = arbiter(inputs)
-        score1 = results[0]['score'] if results[0]['label'].lower() == 'positive' else 1 - results[0]['score']
-        score2 = results[1]['score'] if results[1]['label'].lower() == 'positive' else 1 - results[1]['score']
-        if score1 > score2:
-            return ("GPT-2", resp1)
-        else:
-            return ("GPT-Neo 125M", resp2)
-    except Exception as e:
-        # Em caso de erro, escolha arbitrariamente a primeira
-        return ("GPT-2", resp1)
+    if score_a > score_b:
+        return "Model A", response_a
+    else:
+        return "Model B", response_b
 
+# Função principal do chatbot
 def chatbot(prompt):
-    resp_gpt2 = generate_response(prompt, "GPT-2")
-    resp_gptneo = generate_response(prompt, "GPT-Neo 125M")
-    winner_name, winner_resp = judge_responses(resp_gpt2, resp_gptneo)
-    return {
-        "GPT-2": resp_gpt2,
-        "GPT-Neo 125M": resp_gptneo,
-        "Modelo Vencedor": winner_name,
-        "Resposta Escolhida": winner_resp
-    }
+    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()
 
-with gr.Blocks(title="Chatbot com Árbitro") as demo:
-    gr.Markdown("# 🤖 Chatbot de Comparação com Árbitro")
-    gr.Markdown("Testa dois modelos e escolhe a melhor resposta com um modelo árbitro")
-    
-    with gr.Row():
-        input_prompt = gr.Textbox(
-            label="Digite sua pergunta:",
-            placeholder="Escreva algo em português...",
-            lines=3
-        )
-    
-    submit_btn = gr.Button("Enviar Pergunta", variant="primary")
-    
-    with gr.Row():
-        output_gpt2 = gr.Textbox(label="Resposta GPT-2", interactive=False)
-        output_gptneo = gr.Textbox(label="Resposta GPT-Neo 125M", interactive=False)
+    winner, final_response = judge_response(prompt, response_a, response_b)
 
-    output_winner = gr.Textbox(label="Modelo Vencedor", interactive=False)
-    output_chosen = gr.Textbox(label="Resposta Escolhida", interactive=False)
-    
-    submit_btn.click(
-        fn=chatbot,
-        inputs=input_prompt,
-        outputs=[output_gpt2, output_gptneo, output_winner, output_chosen]
+    return (
+        prompt,
+        response_a,
+        response_b,
+        winner,
+        final_response
     )
 
-def cleanup():
-    global loaded_models
-    for model in loaded_models.values():
-        if model:
-            del model["pipe"]
-    torch.cuda.empty_cache()
+# Interface Gradio
+iface = gr.Interface(
+    fn=chatbot,
+    inputs=gr.Textbox(label="Digite sua pergunta:"),
+    outputs=[
+        gr.Textbox(label="Pergunta"),
+        gr.Textbox(label="Resposta do Modelo A (GPT-2)"),
+        gr.Textbox(label="Resposta do Modelo B (GPT-Neo-125M)"),
+        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.
 
-import atexit
-atexit.register(cleanup)
+    🧠 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`)
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    Critério de julgamento: Resposta com maior sentimento positivo vence.
+    """
+)
 
+iface.launch()