51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
|
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}]
|