22 lines
917 B
Python
22 lines
917 B
Python
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
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"
|
|
)
|
|
|
|
def __call__(self, data):
|
|
inputs = data.pop("inputs", data)
|
|
params = data.pop("parameters", {})
|
|
encoded = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
|
|
output = self.model.generate(
|
|
**encoded,
|
|
max_new_tokens=params.get("max_new_tokens", 600),
|
|
temperature=params.get("temperature", 0.1),
|
|
repetition_penalty=params.get("repetition_penalty", 1.1),
|
|
do_sample=True,
|
|
)
|
|
return self.tokenizer.decode(output[0][encoded["input_ids"].shape[1]:], skip_special_tokens=True) |