File size: 844 Bytes
4d8ef25 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModel
import torch
class EndpointHandler:
def __init__(self, path=""):
# load model and processor from path
self.tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
self.model = AutoModel.from_pretrained(path, trust_remote_code=True).half().cuda()
def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
"""
Args:
data (:dict:):
The payload with the text prompt and generation parameters.
"""
# process input
inputs = data.pop("inputs", data)
history = data.pop("history", None)
response, new_history = self.model.chat(self.tokenizer, inputs, history)
return [{"generated_text": response, "history": new_history}] |