Update handler.py

handler.py

@@ -1,53 +1,28 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from typing import Dict, List, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 class EndpointHandler:
-    def __init__(self, path="d-s-b/meme"):
-        self.tokenizer = AutoTokenizer.from_pretrained("d-s-b/meme")
+    def __init__(self, path=""):
+        # Load from local path, not from hub
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
-            "d-s-b/meme", 
-            torch_dtype="auto", 
+            path,
+            torch_dtype=torch.float16,
             device_map="auto"
         )
-        
-        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.
-
-### Instruction:
-You are a meme expert who knows how to map real-life situations to the correct meme name.
-Please identify the meme name that best fits the given examples_list.
-
-### Input (examples_list):
-{}
-
-### Response:
-"""
-        
+    
     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")
+        inputs = data.get("inputs", "")
+        parameters = data.get("parameters", {})
         
+        encoded = self.tokenizer(inputs, 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
+            **encoded,
+            max_length=parameters.get("max_length", 100),
+            temperature=parameters.get("temperature", 0.7),
+            do_sample=True
         )
         
-        response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        result = response[0].split("### Response:")[1].strip()
-        
-        return [{"generated_text": result}]
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return [{"generated_text": response}]