Create handler.py

handler.py

@@ -0,0 +1,17 @@
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+class EndpointHandler:
+    def __init__(self, model_dir: str, **kw):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+        with init_empty_weights():
+            model = AutoModelForCausalLM.from_pretrained(
+                model_dir, torch_dtype="auto", trust_remote_code=True
+            )
+        self.model = load_checkpoint_and_dispatch(
+            model, checkpoint=model_dir, device_map="auto"
+        )                      # 自动跨 GPU 切层
+    def __call__(self, data):
+        prompt = data["inputs"]
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        out_ids = self.model.generate(**inputs, max_new_tokens=256)
+        return {"generated_text": self.tokenizer.decode(out_ids[0], skip_special_tokens=True)}