Update handler.py

handler.py

@@ -1,77 +1,44 @@
-from typing import Dict, Any, List
+from typing import Dict, Any
 import torch
-from transformers import pipeline, AutoTokenizer, LlamaForCausalLM
-
-dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        self.tokenizer = AutoTokenizer.from_pretrained(path, revision="main")
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-
-        self.model = LlamaForCausalLM.from_pretrained(path, revision="main", torch_dtype=dtype)
-
-        device = 0 if torch.cuda.is_available() else -1
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
+        
         self.generator = pipeline(
             "text-generation",
             model=self.model,
             tokenizer=self.tokenizer,
-            device=device
+            device=0 if torch.cuda.is_available() else -1,
+            return_full_text=False,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16
         )
 
-        self.eos_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.eos_token)
-
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        # ✅ If using HF Inference Endpoint, wrap everything under "inputs"
-        data = data.get("inputs", data)
-
-        messages: List[Dict[str, str]] = data.get("messages", [])
-        if not messages:
-            return {"error": "Missing 'messages' array."}
-
-        prompt = self.format_chat_prompt(messages)
-
-        generation_args = data.get("parameters", {})
-        max_tokens = generation_args.setdefault("max_new_tokens", 300)
-        generation_args.setdefault("do_sample", True)
-        generation_args.setdefault("temperature", 0.4)
-        generation_args.setdefault("top_p", 0.9)
-        generation_args.setdefault("repetition_penalty", 1.2)
-        generation_args.setdefault("no_repeat_ngram_size", 6)
-        generation_args.setdefault("early_stopping", True)
-        generation_args.setdefault("return_full_text", False)
-        generation_args.setdefault("eos_token_id", self.eos_token_id)
-        generation_args.setdefault("pad_token_id", self.tokenizer.pad_token_id)
-
-        try:
-            result = self.generator(prompt, **generation_args)
-            output = result[0]["generated_text"].strip()
-            token_count = len(self.tokenizer.encode(output))
-
-            finish_reason = "stop"
-            if self.tokenizer.eos_token not in output and token_count >= max_tokens:
-                finish_reason = "length"
-
-            return {
-                "choices": [{
-                    "message": {
-                        "role": "assistant",
-                        "content": output
-                    },
-                    "finish_reason": finish_reason
-                }]
-            }
-
-        except Exception as e:
-            import traceback
-            return {"error": str(e), "traceback": traceback.format_exc()}
+        prompt = data.get("inputs", "")
+        if not prompt:
+            return {"error": "Missing 'inputs' field."}
+
+        if not prompt.startswith("<|begin_of_text|>"):
+            prompt = f"<|begin_of_text|>{prompt}"
+
+        params = data.get("parameters", {})
+        outputs = self.generator(
+            prompt,
+            max_new_tokens=params.get("max_new_tokens", 100),
+            do_sample=params.get("do_sample", True),
+            temperature=params.get("temperature", 0.7),
+            top_p=params.get("top_p", 0.9)
+        )
 
-    def format_chat_prompt(self, messages: List[Dict[str, str]]) -> str:
-        prompt = ""
-        for msg in messages:
-            role = msg.get("role", "").strip().lower()
-            content = msg.get("content", "").strip()
-            if role in ["system", "user", "assistant", "ipython"]:
-                prompt += f"{content}\n"
-        return prompt.strip()
+        return {
+            "choices": [{
+                "message": {
+                    "role": "assistant",
+                    "content": outputs[0]["generated_text"].strip()
+                },
+                "finish_reason": "stop"
+            }]
+        }