Create handler.py

handler.py

@@ -0,0 +1,77 @@
+from typing import Dict, Any, List
+import torch
+from transformers import pipeline, AutoTokenizer, LlamaForCausalLM
+
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
+
+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.generator = pipeline(
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            device=device
+        )
+
+        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()}
+
+    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()