adjust model

configuration/config

@@ -1,2 +0,0 @@
-MODEL_PATH_BRISA_7B = "./models/brisa/BRisa-7B-Instruct-v0.2.Q4_K_S.gguf"
-MODEL_PATH_META_LLM_8B = "./models/llama/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"

configuration/config.py

@@ -0,0 +1,2 @@
+MODEL_FLAN_T5 = "google/flan-t5-small"
+MODEL_FALCON_RW_1B = "tiiuae/falcon-rw-1b"

models/model.py

@@ -1,10 +1,25 @@
-from llama_cpp import Llama
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
 
 class Model:
-    def __init__(self, model_path):
-        self.model = Llama(model_path=model_path)
+    def __init__(self, model_name):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-    def generate_response(self, prompt: str) -> str:
-        output = self.model(prompt, max_tokens=60)
-        return output['choices'][0]['text'].strip()
+    def generate_response_model_1(self, question: str) -> str:
+        prompt = f"Give a short and factual answer: {question}"
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_length=30,
+                temperature=0.4,
+                top_k=40,
+                top_p=0.85,
+                repetition_penalty=1.2,
+                do_sample=True
+            )
+
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

service/chatbot.py

@@ -1,23 +1,17 @@
 from models.model import Model
-from config import MODEL_PATH_BRISA_7B, MODEL_PATH_META_LLM_8B
+from configuration.config import MODEL_FLAN_T5, MODEL_FALCON_RW_1B
 from models.arbitrator import Arbitrator
+from sentence_transformers import SentenceTransformer, util
 
-model_a = Model(model_path=MODEL_PATH_BRISA_7B)
-model_b = Model(model_path=MODEL_PATH_META_LLM_8B)
+model_a = Model(model_path=MODEL_FLAN_T5)
+model_b = Model(model_path=MODEL_FALCON_RW_1B)
 arbitrator = Arbitrator()
 
+model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-def process_prompt(prompt: str) -> str:
-    print("Gerando resposta A...")
-    response_a = model_a.generate_response(prompt.strip())
-    print("Resposta A gerada:", response_a)
 
-    print("Gerando resposta B...")
-    response_b = model_b.generate_response(prompt.strip())
-    print("Resposta B gerada:", response_b)
-
-    print("Avaliando...")
-    best_response = arbitrator.evaluate(prompt.strip(), response_a.strip(), response_b.strip())
-    print("Avaliação concluída. Melhor resposta selecionada.")
-
-    return best_response
+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