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

Model: mariadelcarmenramirez/Llama-3.1-8B-ArtTherapy
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-08-05 06:47:16 +08:00
commit 547eae25f8
9 changed files with 425 additions and 0 deletions

51
handler.py Normal file
View File

@@ -0,0 +1,51 @@
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
class EndpointHandler:
def __init__(self, path="/repository"):
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.model = AutoModelForCausalLM.from_pretrained(
path,
torch_dtype=torch.bfloat16,
device_map="auto",
)
self.model.eval()
def __call__(self, data):
inputs = data.get("inputs", data)
parameters = data.get("parameters", {})
max_new_tokens = parameters.get("max_new_tokens", 256)
temperature = parameters.get("temperature", 0.7)
if isinstance(inputs, dict) and "messages" in inputs:
messages = inputs["messages"]
# Step 1: apply template to get a STRING (tokenize=False)
prompt = self.tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
# Step 2: tokenize the string
input_ids = self.tokenizer.encode(
prompt, return_tensors="pt"
).to(self.model.device)
else:
text = inputs if isinstance(inputs, str) else str(inputs)
input_ids = self.tokenizer.encode(
text, return_tensors="pt"
).to(self.model.device)
with torch.no_grad():
output = self.model.generate(
input_ids,
max_new_tokens=max_new_tokens,
temperature=temperature,
do_sample=True,
)
result = self.tokenizer.decode(
output[0][input_ids.shape[-1]:], skip_special_tokens=True
)
return [{"generated_text": result}]