From 0f1c2885305259b8f00b0f53c247964c1cd80d64 Mon Sep 17 00:00:00 2001 From: Wen Date: Wed, 29 Mar 2023 14:21:13 +0000 Subject: [PATCH] Create handler.py --- handler.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 handler.py diff --git a/handler.py b/handler.py new file mode 100644 index 0000000..9e0cfdc --- /dev/null +++ b/handler.py @@ -0,0 +1,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}] \ No newline at end of file