Create handler.py

handler.py

@@ -0,0 +1,48 @@
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+)
+import torch, os
+
+MODEL_ID = "Qwen/Qwen3-8B-Instruct"   # 换成自己的模型
+
+def get_model():
+    # ① 先试 bfloat16，A100/H100 都原生支持
+    return AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",          # TGI 同款逻辑，自动分片 
+        low_cpu_mem_usage=True,     # 先在 CPU 建图，再流式拷到 GPU
+        trust_remote_code=True
+    )
+
+# ---- 如果 bfloat16 仍 OOM，可改成 4-bit 量化 ----
+# bnb_cfg = BitsAndBytesConfig(
+#     load_in_4bit=True,
+#     bnb_4bit_quant_type="nf4",
+#     bnb_4bit_use_double_quant=True,
+# )
+# def get_model():
+#     return AutoModelForCausalLM.from_pretrained(
+#         MODEL_ID,
+#         device_map="auto",
+#         quantization_config=bnb_cfg,
+#         trust_remote_code=True
+#     )
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = get_model()
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto",
+    torch_dtype=getattr(model, "dtype", torch.bfloat16),
+)
+
+def __init__(self, *args, **kwargs):
+    pass
+
+def __call__(self, data):
+    prompt = data.get("inputs") if isinstance(data, dict) else data
+    outputs = generator(prompt, max_new_tokens=256)
+    return outputs