File size: 1,402 Bytes
9e7f6c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | from transformers import AutoModelForCausalLM, AutoTokenizer
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.model.eval()
def __call__(self, data):
inputs = data.pop("inputs", "")
parameters = data.pop("parameters", {})
max_new_tokens = parameters.get("max_new_tokens", 128)
temperature = parameters.get("temperature", 0.7)
top_p = parameters.get("top_p", 0.9)
repetition_penalty = parameters.get("repetition_penalty", 1.1)
tokenized = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
with torch.no_grad():
outputs = self.model.generate(
**tokenized,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
repetition_penalty=repetition_penalty,
do_sample=temperature > 0,
)
# Decode only the NEW tokens (exclude the input prompt)
new_tokens = outputs[0][tokenized["input_ids"].shape[1]:]
result = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
return [{"generated_text": result}]
|