| 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) |
