Update model_keeper.py

model_keeper.py

@@ -1,4 +1,4 @@
-from configuration_keeper import KeeperConfig
+from .configuration_keeper import KeeperConfig
 
 import torch
 from transformers import (
@@ -61,45 +61,64 @@ class KeeperModelForCausalLM(PreTrainedModel):
         self.update_both = update_both
         print(f"Model n_cands: {self.n_cands}")
 
-        # Inicializar buffers vacíos para document_vecs y document_mask
-        self.register_buffer('document_retriever_text', torch.empty(0, dtype=torch.long))
-        self.register_buffer('document_retriever_mask', torch.empty(0, dtype=torch.long))
-        self.register_buffer('document_retriever_type', torch.empty(0, dtype=torch.long))
-        self.register_buffer('document_model_text', torch.empty(0, dtype=torch.long))
-        # self.register_buffer('document_model_mask', torch.empty(0, dtype=torch.long))
-        # self.register_buffer('document_model_type', torch.empty(0, dtype=torch.long))
-        self.register_buffer('prompt_left', torch.empty(0, dtype=torch.long))
-        self.register_buffer('prompt_right', torch.empty(0, dtype=torch.long))
-        self.register_buffer('respuesta', torch.empty(0, dtype=torch.long))
 
-    def generate(self, query: Dict[str, torch.LongTensor], k: int = 3, **kwargs):
+    def _load_from_state_dict(self, state_dict, *args, **kwargs):
+        super()._load_from_state_dict(state_dict, *args, **kwargs)
+        # Ensure CUDA is available
+        if torch.cuda.is_available():
+            device = torch.device('cuda')
+            if "document_retriever_text" in state_dict:
+                self.document_retriever_text = state_dict["document_retriever_text"].to(device)
+            if "document_retriever_mask" in state_dict:
+                self.document_retriever_mask = state_dict["document_retriever_mask"].to(device)
+            if "document_retriever_type" in state_dict:
+                self.document_retriever_type = state_dict["document_retriever_type"].to(device)
+            if "document_model_text" in state_dict:
+                self.document_model_text = state_dict["document_model_text"].to(device)
+            if "prompt_left" in state_dict:
+                self.prompt_left = state_dict["prompt_left"].to(device)
+            if "prompt_right" in state_dict:
+                self.prompt_right = state_dict["prompt_right"].to(device)
+            if "respuesta" in state_dict:
+                self.respuesta = state_dict["respuesta"].to(device)
+        else:
+            # Optionally handle the case where CUDA is not available
+            print("CUDA is not available. Tensors will remain on CPU.")
+
+
+    def generate(self, query: Dict[str, torch.LongTensor], k: int = 3,  max_new_tokens=256, repetition_penalty=1.15, temperature=0.1, do_sample=True, **kwargs):
 
-        query_retriever = {k: v.to("cuda") for k, v in query['tokens_retriever'].items()}
         query_model = {k: v.to("cuda") for k, v in query['tokens_model'].items()}
 
-        query_vecs = self.forward_representation(query_retriever)
+        topk_texts = self.document_extractor(query, k)
 
-        doc_dic = {'input_ids': self.document_retriever_text, 'attention_mask':self.document_retriever_mask, 'token_type_ids': self.document_retriever_type}
+        concatenated_texts = torch.cat(topk_texts, dim=0)
 
-        document_vecs = self.forward_representation(doc_dic, sequence_type="doc")
+        T = torch.cat((self.prompt_left, concatenated_texts.unsqueeze(0), self.prompt_right, query_model['input_ids'], self.respuesta), dim=1)
 
-        self.score = self.forward_aggregation(query_vecs, query['tokens_model']["attention_mask"], document_vecs, self.document_retriever_mask)
+        prompt_length = T.shape[1]
 
-        k = min(k, self.score.numel())
+        outputs = self.llm.generate(input_ids=T,max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty, temperature=temperature, do_sample=do_sample)
 
-        topk_scores, topk_indices = torch.topk(self.score, k)
+        return outputs[0][prompt_length:].unsqueeze(0)
 
-        topk_texts = [self.document_model_text[i] for i in topk_indices[0].tolist()]
+    def document_extractor(self, query: Dict[str, torch.LongTensor], k_val: int = 3, **kwargs):
 
-        concatenated_texts = torch.cat(topk_texts, dim=0)
+        query_retriever = {k: v.to("cuda") for k, v in query['tokens_retriever'].items()}
 
-        T = torch.cat((self.prompt_left, concatenated_texts.unsqueeze(0), self.prompt_right, query_model['input_ids'], self.respuesta), dim=1)
+        query_vecs = self.forward_representation(query_retriever)
 
-        prompt_length = T.shape[1]
+        doc_dic = {'input_ids': self.document_retriever_text, 'attention_mask':self.document_retriever_mask, 'token_type_ids': self.document_retriever_type}
+
+        document_vecs = self.forward_representation(doc_dic, sequence_type="doc")
 
-        outputs = self.llm.generate(input_ids=T, max_new_tokens=256, repetition_penalty=1.15)
+        self.score = self.forward_aggregation(query_vecs, query['tokens_retriever']["attention_mask"], document_vecs, self.document_retriever_mask)
 
-        return outputs[0][prompt_length:].unsqueeze(0)
+        k_val = min(k_val, self.score.numel())
+
+        topk_scores, topk_indices = torch.topk(self.score, k_val)
+
+        return [self.document_model_text[i,:] for i in topk_indices[0].tolist()]
 
     def forward_representation(self,
                                tokens,
@@ -145,16 +164,16 @@ class KeeperModelForCausalLM(PreTrainedModel):
     def prompt(self, left_p = None, right_p = None):
         if left_p is None:
           left_p = """ <bos><start_of_turn>user 
-          Eres un experto en cultura paraguaya que responde segun el contexto:
+          Eres un experto en cultura paraguaya que responde de forma clara, amable y concisa. 
+          Segun el siguiente contexto:
 -------------------------------
 """
         if right_p is None:
           right_p = """
 -------------------------------
-- Debes responder solamente en Espanol
-- No utilices conocimientos previos.
-- Responde de forma clara, amable y concisa.
-
+- Solamente puedes responder usando el contexto de arriba, si no se encuentra en el contexto mencionar: 'No tengo informacion sobre eso'.
+- Si encuentras la respuesta puedes copiarla.
+- Debes responder solamente en Espanol.
 Pregunta: """
         return left_p, right_p
 
@@ -187,5 +206,4 @@ Respuesta: <end_of_turn>
         # self.document_model_type = key_outputs['tokens_model']['token_type_ids']
         self.prompt_left = prompt_left_output['tokens_model']['input_ids']
         self.prompt_right = prompt_right_output['tokens_model']['input_ids']
-        self.respuesta = resp_model['input_ids']
-
+        self.respuesta = resp_model['input_ids']