Update handler.py

handler.py

@@ -1,13 +1,18 @@
-from typing import Dict, List, Any
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from typing import Dict, List, Any
 
 class EndpointHandler:
-   def __init__(self, path=""):
-       from transformers import AutoModelForCausalLM, AutoTokenizer
-       self.tokenizer = AutoTokenizer.from_pretrained(path)
-       self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype="auto", device_map="auto")
-       
-       self.inference_prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
+    def __init__(self, path="d-s-b/meme"):
+        self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            path, 
+            torch_dtype="auto", 
+            device_map="auto",
+            trust_remote_code=True
+        )
+        
+        self.inference_prompt_style = """Below is an instruction that describes a task, paired with an input that provides further context.
 Write a response that appropriately completes the request.
 Identify the most suitable meme template based on the provided example situations.
 
@@ -20,32 +25,30 @@ Please identify the meme name that best fits the given examples_list.
 
 ### Response:
 """
-   
-   def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-       question = data.pop("inputs", data)
-       parameters = data.pop("parameters", {})
-       
-       # Set default parameters
-       max_new_tokens = parameters.get("max_new_tokens", 12000)
-       
-       # Format prompt
-       prompt = self.inference_prompt_style.format(question) + self.tokenizer.eos_token
-       
-       # Tokenize
-       inputs = self.tokenizer([prompt], return_tensors="pt")
-       
-       # Generate
-       outputs = self.model.generate(
-           input_ids=inputs.input_ids,
-           attention_mask=inputs.attention_mask,
-           max_new_tokens=max_new_tokens,
-           eos_token_id=self.tokenizer.eos_token_id,
-           use_cache=True,
-           **parameters
-       )
-       
-       # Decode and extract response
-       response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-       result = response[0].split("### Response:")[1]
-       
-       return [{"generated_text": result}]
+        
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        question = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
+        
+        max_new_tokens = parameters.get("max_new_tokens", 512)
+        
+        prompt = self.inference_prompt_style.format(question)
+        
+        inputs = self.tokenizer([prompt], return_tensors="pt")
+        
+        outputs = self.model.generate(
+            input_ids=inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=max_new_tokens,
+            temperature=0.7,
+            do_sample=True,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.eos_token_id,
+            use_cache=True,
+            **parameters
+        )
+        
+        response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        result = response[0].split("### Response:")[1].strip()
+        
+        return [{"generated_text": result}]