| from typing import Dict, List, Any | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| torch.set_default_device("cuda") | |
| self.model = AutoModelForCausalLM.from_pretrained("chuckfinca/arithmephi", torch_dtype="auto", trust_remote_code=True, device_map = 'cuda') | |
| self.tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True) | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| input = data.get("inputs", data) | |
| inputs = self.tokenizer(input, return_tensors="pt", return_attention_mask=False).to('cuda') | |
| outputs = self.model.generate(**inputs, max_length=len(inputs[0]) + 8, pad_token_id=self.tokenizer.eos_token_id) | |
| text = self.tokenizer.batch_decode(outputs)[0] | |
| return text |