wang-sy commited on
Commit
845c287
·
1 Parent(s): 77d4988

Upload 2 files

Browse files

added requirements.txt and handler.py

Files changed (2) hide show
  1. handler.py +22 -0
  2. requirements.txt +5 -0
handler.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import pipeline, GPT2Tokenizer
3
+ from model import GPT
4
+
5
+ class EndpointHandler():
6
+ def __init__(self, path=""):
7
+ # Preload all the elements you are going to need at inference.
8
+ model = GPT.from_pretrained(path)
9
+ tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
10
+ self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
11
+
12
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
13
+ """
14
+ data args:
15
+ inputs (:obj: `str` | `PIL.Image` | `np.array`)
16
+ kwargs
17
+ Return:
18
+ A :obj:`list` | `dict`: will be serialized and returned
19
+ """
20
+ inputs = data.pop("inputs", data)
21
+ output = self.pipeline(inputs)
22
+ return {"Answer": output}
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu116
2
+ torch
3
+ tiktoken
4
+ numpy
5
+ transformers