Update handler.py

handler.py

@@ -1,29 +1,64 @@
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+# handler.py  ——  放在模型仓库根目录
+from typing import Dict, Any
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+
+
 class EndpointHandler:
-    def __init__(self, model_dir: str, **kw):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+    """
+    Hugging Face Inference Endpoints 约定的自定义入口：
+      • __init__(model_dir, **kwargs)   —— 加载模型
+      • __call__(inputs: Dict) -> Dict  —— 处理一次请求
+    """
+
+    def __init__(self, model_dir: str, **kwargs):
+        # 1️⃣ Tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_dir, trust_remote_code=True
+        )
+
+        # 2️⃣ 构建“空壳”模型（不占显存）
         with init_empty_weights():
-            model = AutoModelForCausalLM.from_pretrained(
-                model_dir, torch_dtype="auto", trust_remote_code=True
+            base_model = AutoModelForCausalLM.from_pretrained(
+                model_dir,
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
             )
+
+        # 3️⃣ 把权重切片加载到两张 GPU
         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("cuda:0")          # 👈 把 input_ids/attention_mask 都放到 0 号卡
-
-    out_ids = self.model.generate(
-        **inputs,
-        max_new_tokens=256,
-    )
-    return {
-        "generated_text": self.tokenizer.decode(
-            out_ids[0], skip_special_tokens=True
+            base_model,
+            checkpoint=model_dir,
+            device_map="auto",                # 自动分层到 cuda:0 / cuda:1
+            dtype=torch.float16,
+        )
+
+        # 4️⃣ 生成时常用的生成参数
+        self.generation_kwargs = dict(
+            max_new_tokens=2048,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
         )
-    }
 
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
+        """
+        data 格式：
+          {
+            "inputs": "your prompt here"
+          }
+        """
+        prompt = data["inputs"]
+
+        # ➡️ 只把输入张量放到 cuda:0（与模型第一层同卡）
+        inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda:0")
+
+        # 生成
+        with torch.inference_mode():
+            output_ids = self.model.generate(**inputs, **self.generation_kwargs)
+
+        generated_text = self.tokenizer.decode(
+            output_ids[0], skip_special_tokens=True
+        )
+        return {"generated_text": generated_text}