| import torch | |
| from typing import Dict, List, Any | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| MAX_TOKENS_IN_BATCH = 4_000 | |
| DEFAULT_MAX_NEW_TOKENS = 10 | |
| class EndpointHandler(): | |
| def __init__(self, path: str = ""): | |
| self.tokenizer = AutoTokenizer.from_pretrained(path) | |
| self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16) | |
| self.model = self.model.to('cuda:0') | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| """ | |
| Args: | |
| data (:obj:): | |
| includes the input data and the parameters for the inference. | |
| Return: | |
| A :obj:`list`:. The list contains the answer and scores of the inference inputs | |
| """ | |
| prompts = [f"<human>: {prompt}\n<bot>:" for prompt in data["inputs"]] | |
| self.tokenizer.pad_token = self.tokenizer.eos_token | |
| inputs = self.tokenizer(prompts, padding=True, return_tensors='pt').to(self.model.device) | |
| input_length = inputs.input_ids.shape[1] | |
| outputs = self.model.generate( | |
| **inputs, **data["parameters"] | |
| ) | |
| output_strs = self.tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True) | |
| return [{"generated_text": output_strs}] |