Update handler.py

handler.py

@@ -4,7 +4,6 @@ import torch
 import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-
 dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
 
 class EndpointHandler:
@@ -18,45 +17,27 @@ class EndpointHandler:
             torch_dtype=dtype,
             trust_remote_code=True,
         )
-               
+        
         generation_config = self.model.generation_config
         generation_config.max_new_tokens = 200
-        generation_config.temperature = 0.8
+        generation_config.temperature = 0.4
         generation_config.top_p = 0.8
         generation_config.num_return_sequences = 1
         generation_config.pad_token_id = self.tokenizer.eos_token_id
         generation_config.eos_token_id = self.tokenizer.eos_token_id
-        generation_config.early_stopping = True
         self.generate_config = generation_config
         
         self.pipeline = transformers.pipeline(
             "text-generation", model=self.model, tokenizer=self.tokenizer
         )
 
-
-    def _ensure_token_limit(self, text):
-        """Ensure text is within the model's token limit."""
-        tokens = self.tokenizer.tokenize(text)
-        if len(tokens) > 2048:
-            # Remove tokens from the beginning until the text fits
-            tokens = tokens[-2048:]
-            return self.tokenizer.decode(tokens)
-        return text
-    
-    
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         user_prompt = data.pop("inputs", data)
         
-        # Permanent context
-        permanent_context = "<context>: You are a life coaching bot..."
-        structured_prompt = f"{permanent_context}\<bot> response:"
-        
-        result = self.pipeline(structured_prompt, generation_config=self.generate_config)
+        # Add the permanent context to the user's prompt
+        permanent_context = "<context>: You are a life coaching bot with the goal of improving understanding, reducing suffering and improving life. Learn about the user in order to provide guidance without making assumptions or adding information not provided by the user."
+        combined_prompt = f"{permanent_context}\n<human>: {user_prompt}"
         
-        # Ensure _extract_response is defined and works as intended
-        response_text = self._extract_response(result[0]['generated_text'])
+        result = self.pipeline(combined_prompt, generation_config=self.generate_config)
         
-        # Trimming response
-        response_text = response_text.rsplit("[END", 1)[0].strip()
-    
-        return {"response": response_text}
+        return result