初始化项目,由ModelHub XC社区提供模型

Model: Prod5/mistral-7b-a2ui
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-06-02 02:06:21 +08:00
commit 7d10129465
13 changed files with 105513 additions and 0 deletions

27
handler.py Normal file
View File

@@ -0,0 +1,27 @@
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
class EndpointHandler:
def __init__(self, path=""):
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(
path,
torch_dtype=torch.float16,
device_map="auto",
)
self.pipeline = pipeline(
"text-generation",
model=self.model,
tokenizer=self.tokenizer,
)
def __call__(self, data):
inputs = data.get("inputs", "")
parameters = data.get("parameters", {})
if not parameters.get("max_new_tokens"):
parameters["max_new_tokens"] = 2048
result = self.pipeline(inputs, **parameters)
return result