Upload 3 files

handler.py

@@ -1,68 +1,45 @@
-import os
+from typing import Dict, List, Any
+import json
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
-# 定义模型处理类
-class ModelHandler(object):
-    def __init__(self):
+class EndpointHandler():
+    def __init__(self, path=""):
         self.tokenizer = None
         self.model = None
         self.device = None
+        self.load_model(path)
 
     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.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).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:
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", data)
+        print(f'get input {inputs}')
+        if not inputs:
             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})
+        history = json.loads(inputs)
+        print(f'history is {history}')
+        #history.append({"role": "user", "content": inputs})
         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):
+        print(f'encoding success')
         with torch.no_grad():
             output = self.model.generate(
-                **model_input,
+                **encoding,
                 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)
-
+        print(f'output success')
+        generated_text = self.tokenizer.decode(output[0], skip_special_tokens=True)
 
-def postprocess(prediction):
-    return _service.postprocess(prediction)
+        return [{"response": generated_text}]

modeling_minicpm.py

@@ -38,7 +38,7 @@ from transformers.modeling_attn_mask_utils import (
 )
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
-from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_13
+from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
 from transformers.utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
@@ -61,8 +61,8 @@ except:
 # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
 # It means that the function will not be traced through and simply appear as a node in the graph.
 if is_torch_fx_available():
-    if not is_torch_greater_or_equal_than_1_13:
-        import torch.fx
+    # if not is_torch_greater_or_equal_than_1_13:
+    #     import torch.fx
 
     _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
 

requirements.txt

@@ -0,0 +1,5 @@
+# for MiniCPM-2B hf inference
+torch>=2.0.0
+transformers>=4.36.2
+
+