Upload handler.py

handler.py

@@ -0,0 +1,68 @@
+import os
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+
+# 定义模型处理类
+class ModelHandler(object):
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+        self.device = None
+
+    def load_model(self, model_dir):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        model_path = model_dir
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path).to(self.device)
+        self.model.eval()
+
+        print(f"Tokenizer and Model loaded from: {model_path} to device: {self.device}")
+
+
+    def preprocess(self, request):
+        input_text = request.get("inputs", request.get("text"))
+        if not input_text:
+            raise ValueError("Input text is missing in the request. Please provide 'inputs' or 'text' in your request.")
+
+        history = []
+        history.append({"role": "user", "content": input_text})
+        conversion = self.tokenizer.apply_chat_template(history, add_generation_prompt=True, tokenize=False)
+        encoding = self.tokenizer(conversion, return_tensors="pt").to(self.device)
+        return encoding
+
+
+    def predict(self, model_input):
+        with torch.no_grad():
+            output = self.model.generate(
+                **model_input,
+                max_new_tokens=1024,
+                temperature=1.5,
+                do_sample=True,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+        return output
+
+
+    def postprocess(self, prediction):
+        generated_text = self.tokenizer.decode(prediction[0], skip_special_tokens=True)
+        return {"response": generated_text}
+
+
+_service = ModelHandler()
+
+def load():
+    model_dir = '/home/aistudio/export'
+    _service.load_model(model_dir)
+
+
+def preprocess(request):
+    return _service.preprocess(request)
+
+
+def predict(data):
+    return _service.predict(data)
+
+
+def postprocess(prediction):
+    return _service.postprocess(prediction)