fwittel
commited on
Commit
•
ba13b5a
1
Parent(s):
97ce3d1
added custom handler
Browse files- __pycache__/handler.cpython-311.pyc +0 -0
- handler.py +16 -0
- requirements.txt +1 -0
- test.py +13 -0
__pycache__/handler.cpython-311.pyc
ADDED
Binary file (1.38 kB). View file
|
|
handler.py
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Dict, List, Any
|
2 |
+
from transformers import AutoModel, pipeline
|
3 |
+
|
4 |
+
class EndpointHandler:
|
5 |
+
def __init__(self, path=""):
|
6 |
+
# load the model
|
7 |
+
model = AutoModel.from_pretrained(path, low_cpu_mem_usage=True)
|
8 |
+
# create inference pipeline
|
9 |
+
# Do I have to check device?
|
10 |
+
self.pipeline = pipeline("text-generation", model=model)
|
11 |
+
|
12 |
+
# (Might have to adjust typing)
|
13 |
+
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
|
14 |
+
inputs = data.pop("inputs", data) # Should I get and pass parameters?
|
15 |
+
prediction = self.pipeline(inputs)
|
16 |
+
return prediction
|
requirements.txt
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
transformers
|
test.py
ADDED
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from handler import EndpointHandler
|
2 |
+
|
3 |
+
# init handler
|
4 |
+
my_handler = EndpointHandler(path=".")
|
5 |
+
|
6 |
+
# prepare sample payload
|
7 |
+
payload = {"inputs": "I am Bob and I want to "}
|
8 |
+
|
9 |
+
# test the handler
|
10 |
+
pred=my_handler(payload)
|
11 |
+
|
12 |
+
# show results
|
13 |
+
print("pred", pred)
|