fwittel commited on
Commit
ba13b5a
1 Parent(s): 97ce3d1

added custom handler

Browse files
Files changed (4) hide show
  1. __pycache__/handler.cpython-311.pyc +0 -0
  2. handler.py +16 -0
  3. requirements.txt +1 -0
  4. 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)