adjust model

app.py

@@ -1,16 +1,21 @@
 import gradio as gr
-from service.chatbot import process_prompt
+from logic.pipeline import process_question
 
 
-def main():
-    interface = gr.Interface(
-        fn=process_prompt,
-        inputs=gr.Textbox(label="Pergunta do Usuário"),
-        outputs=gr.Textbox(label="Melhor Resposta"),
-        title="🤖 Chatbot em Cascata com Avaliação Rápida"
-    )
-    interface.launch()
+def handle_input(question):
+    result = process_question(question)
+    return result["answer_model_1"], result["answer_model_2"], result["best_answer"]
 
 
-if __name__ == '__main__':
-    main()
+iface = gr.Interface(
+    fn=handle_input,
+    inputs=gr.Textbox(label="Pergunta"),
+    outputs=[
+        gr.Textbox(label="Resposta do Modelo 1"),
+        gr.Textbox(label="Resposta do Modelo 2"),
+        gr.Textbox(label="Melhor Resposta (segundo Avaliador)")
+    ]
+)
+
+if __name__ == "__main__":
+    iface.launch()

models/arbitrator.py

@@ -1,16 +1,10 @@
 from sentence_transformers import SentenceTransformer, util
 
+model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-class Arbitrator:
-    def __init__(self):
-        self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
 
-    def evaluate(self, prompt: str, response_a: str, response_b: str) -> str:
-        prompt_emb = self.model.encode(prompt.strip(), convert_to_tensor=True)
-        emb_a = self.model.encode(response_a.strip(), convert_to_tensor=True)
-        emb_b = self.model.encode(response_b.strip(), convert_to_tensor=True)
-
-        sim_a = util.cos_sim(prompt_emb, emb_a)
-        sim_b = util.cos_sim(prompt_emb, emb_b)
-
-        return response_a if sim_a > sim_b else response_b
+def choose_best_response(question, answer1, answer2):
+    embeddings = model.encode([question, answer1, answer2])
+    score1 = util.cos_sim(embeddings[0], embeddings[1])
+    score2 = util.cos_sim(embeddings[0], embeddings[2])
+    return answer1 if score1 > score2 else answer2

models/model.py

@@ -22,4 +22,4 @@ class Model:
                 do_sample=True
             )
 
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

requirements.txt

@@ -1,8 +1,5 @@
 huggingface_hub==0.25.2
-gradio==4.26.0
-transformers>=4.40.0
-torch>=1.13
-sentence-transformers>=2.2.2
-accelerate>=1.7.0
-bitsandbytes>=0.46.0
-sentencepiece>=0.2.0
+transformers
+torch
+gradio
+sentence-transformers

service/chatbot.py

@@ -1,17 +1,18 @@
 from models.model import Model
 from configuration.config import MODEL_FLAN_T5, MODEL_FALCON_RW_1B
-from models.arbitrator import Arbitrator
-from sentence_transformers import SentenceTransformer, util
+from models.arbitrator import choose_best_response
 
-model_a = Model(model_path=MODEL_FLAN_T5)
-model_b = Model(model_path=MODEL_FALCON_RW_1B)
-arbitrator = Arbitrator()
+model_a = Model(model_name=MODEL_FLAN_T5)
+model_b = Model(model_name=MODEL_FALCON_RW_1B)
 
-model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-
-def process_prompt(prompt, model_a, model_b):
-    embeddings = model.encode([prompt, model_a, model_b])
-    score1 = util.cos_sim(embeddings[0], embeddings[1])
-    score2 = util.cos_sim(embeddings[0], embeddings[2])
-    return model_a if score1 > score2 else model_b
+def process_question(question):
+    answer1 = model_a.generate_response_model_1(question)
+    answer2 = model_b.generate_response_model_2(question)
+    best = choose_best_response(question, answer1, answer2)
+    return {
+        "question": question,
+        "answer_model_1": answer1,
+        "answer_model_2": answer2,
+        "best_answer": best
+    }