| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
| import torch |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=False) |
| self.model = AutoModelForCausalLM.from_pretrained( |
| path, torch_dtype=torch.float16, device_map="auto" |
| ) |
| self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer) |
|
|
| def __call__(self, data): |
| messages = data.get("inputs", {}).get("messages", []) |
| if not messages: |
| messages = [{"role": "user", "content": str(data.get("inputs", ""))}] |
| text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| out = self.pipe(text, |
| max_new_tokens=data.get("parameters", {}).get("max_tokens", 2048), |
| temperature=data.get("parameters", {}).get("temperature", 0.3), |
| do_sample=True) |
| return {"choices": [{"message": {"role": "assistant", "content": out[0]["generated_text"][len(text):]}}]} |
|
|
