| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from transformers import GenerationConfig, pipeline | |
| import torch | |
| from typing import Any, Dict | |
| dtype = torch.float16 | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| tokenizer = AutoTokenizer.from_pretrained(path) | |
| model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype) | |
| self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| def __call__(self, data: Dict[str, Any]) -> [str]: | |
| inputs = data.pop("inputs", data) | |
| generation_config = GenerationConfig( | |
| max_new_tokens=250, do_sample=True, top_k=50, | |
| temperature=0.8, pad_token_id=2, num_return_sequences=1, | |
| min_new_tokens=60, repetition_penalty=1.2, return_full_text=False | |
| ) | |
| output = self.pipeline(inputs, **generation_config.to_dict()) | |
| return output |