Update app.py

app.py

@@ -1,27 +1,29 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
-import torch
+from transformers import pipeline
+from sentence_transformers import SentenceTransformer, util
 
-# Modelos de resposta A e B
-generator_a = pipeline("text-generation", model="tiiuae/falcon-rw-1b", tokenizer="tiiuae/falcon-rw-1b", device=-1)
-generator_b = pipeline("text2text-generation", model="google/flan-t5-base", tokenizer="google/flan-t5-base", device=-1)
+# Modelos A e B
+generator_a = pipeline("text2text-generation", model="google/flan-t5-base", tokenizer="google/flan-t5-base")
+generator_b = pipeline("text2text-generation", model="declare-lab/flan-alpaca-base", tokenizer="declare-lab/flan-alpaca-base")
 
-# Modelo árbitro
-tokenizer_c = AutoTokenizer.from_pretrained("facebook/bart-large-mnli")
-model_c = AutoModelForSequenceClassification.from_pretrained("facebook/bart-large-mnli")
-classifier = pipeline("text-classification", model=model_c, tokenizer=tokenizer_c)
+# Árbitro baseado em similaridade semântica
+similarity_model = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L6-v2")
 
 def judge(prompt, response_a, response_b):
-    # Avaliar qual resposta tem mais "entailment" com o prompt
-    result_a = classifier(f"{prompt} </s> {response_a}")[0]
-    result_b = classifier(f"{prompt} </s> {response_b}")[0]
-    return response_a if result_a['score'] > result_b['score'] else response_b
+    emb_prompt = similarity_model.encode(prompt, convert_to_tensor=True)
+    emb_a = similarity_model.encode(response_a, convert_to_tensor=True)
+    emb_b = similarity_model.encode(response_b, convert_to_tensor=True)
+    
+    score_a = util.pytorch_cos_sim(emb_prompt, emb_a).item()
+    score_b = util.pytorch_cos_sim(emb_prompt, emb_b).item()
+    
+    return response_a if score_a >= score_b else response_b
 
 def chat(prompt):
-    response_a = generator_a(prompt, max_length=100, do_sample=True)[0]['generated_text']
-    response_b = generator_b(prompt, max_length=100, do_sample=True)[0]['generated_text']
+    response_a = generator_a(prompt, max_length=100)[0]['generated_text']
+    response_b = generator_b(prompt, max_length=100)[0]['generated_text']
     best = judge(prompt, response_a, response_b)
     return f"Resposta A:\n{response_a}\n\nResposta B:\n{response_b}\n\n✅ Melhor Resposta:\n{best}"
 
-iface = gr.Interface(fn=chat, inputs=gr.Textbox(label="Seu prompt"), outputs=gr.Textbox(label="Resposta escolhida"), title="Chatbot em Cascata")
+iface = gr.Interface(fn=chat, inputs=gr.Textbox(label="Pergunta"), outputs=gr.Textbox(label="Resposta Escolhida"), title="Chatbot em Cascata")
 iface.launch()