64 lines
2.6 KiB
Python
64 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
class EndpointHandler:
|
|
def __init__(self, path: str = ""):
|
|
self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
|
|
if self.tokenizer.pad_token_id is None:
|
|
self.tokenizer.pad_token = self.tokenizer.eos_token
|
|
|
|
model_kwargs: dict[str, Any] = {
|
|
"device_map": "auto",
|
|
"trust_remote_code": True,
|
|
"low_cpu_mem_usage": True,
|
|
"attn_implementation": "sdpa",
|
|
}
|
|
if torch.cuda.is_available():
|
|
model_kwargs["torch_dtype"] = torch.bfloat16
|
|
|
|
self.model = AutoModelForCausalLM.from_pretrained(path, **model_kwargs)
|
|
self.model.eval()
|
|
|
|
def __call__(self, data: dict[str, Any]) -> dict[str, str]:
|
|
inputs = data.get("inputs", "")
|
|
parameters = dict(data.get("parameters") or {})
|
|
|
|
if isinstance(inputs, dict) and isinstance(inputs.get("messages"), list):
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
inputs["messages"],
|
|
tokenize=False,
|
|
add_generation_prompt=True,
|
|
)
|
|
else:
|
|
prompt = str(inputs)
|
|
|
|
encoded = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
|
|
max_input_tokens = int(parameters.pop("max_input_tokens", 4096))
|
|
if encoded["input_ids"].shape[1] > max_input_tokens:
|
|
encoded["input_ids"] = encoded["input_ids"][:, -max_input_tokens:]
|
|
encoded["attention_mask"] = encoded["attention_mask"][:, -max_input_tokens:]
|
|
|
|
input_length = encoded["input_ids"].shape[1]
|
|
temperature = float(parameters.pop("temperature", 0))
|
|
do_sample = bool(parameters.pop("do_sample", temperature > 0))
|
|
generation_kwargs = {
|
|
"max_new_tokens": int(parameters.pop("max_new_tokens", 256)),
|
|
"repetition_penalty": float(parameters.pop("repetition_penalty", 1.1)),
|
|
"do_sample": do_sample,
|
|
"pad_token_id": self.tokenizer.pad_token_id,
|
|
"eos_token_id": self.tokenizer.eos_token_id,
|
|
}
|
|
if do_sample:
|
|
generation_kwargs["temperature"] = max(temperature, 0.01)
|
|
generation_kwargs["top_p"] = float(parameters.pop("top_p", 0.9))
|
|
|
|
with torch.inference_mode():
|
|
output = self.model.generate(**encoded, **generation_kwargs)
|
|
generated = self.tokenizer.decode(output[0, input_length:], skip_special_tokens=True).strip()
|
|
return {"generated_text": generated}
|