philschmid HF staff commited on
Commit
d3ce09e
1 Parent(s): 7497345

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +41 -0
handler.py ADDED
@@ -0,0 +1,41 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from typing import Dict, List, Any
3
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
4
+
5
+ # check for GPU
6
+ device = 0 if torch.cuda.is_available() else -1
7
+
8
+ # multi-model list
9
+ multi_model_list = [
10
+ {"model_id": "distilbert-base-uncased-finetuned-sst-2-english", "task": "text-classification"},
11
+ {"model_id": "Helsinki-NLP/opus-mt-en-de", "task": "translation"},
12
+ {"model_id": "facebook/bart-large-cnn", "task": "summarization"},
13
+ {"model_id": "dslim/bert-base-NER", "task": "token-classification"},
14
+ {"model_id": "textattack/bert-base-uncased-ag-news", "task": "text-classification"},
15
+ ]
16
+
17
+
18
+ class EndpointHandler():
19
+ def __init__(self, path=""):
20
+ self.multi_model={}
21
+ # load all the models onto device
22
+ for model in multi_model_list:
23
+ self.multi_model[model["model_id"]] = pipeline(model["task"], model=model["model_id"], device=device)
24
+
25
+ def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
26
+ # deserialize incomin request
27
+ inputs = data.pop("inputs", data)
28
+ parameters = data.pop("parameters", None)
29
+ model_id = data.pop("model_id", None)
30
+
31
+ # check if model_id is in the list of models
32
+ if model_id is None or model_id not in self.multi_model:
33
+ raise ValueError(f"model_id: {model_id} is not valid. Available models are: {list(self.multi_model.keys())}")
34
+
35
+ # pass inputs with all kwargs in data
36
+ if parameters is not None:
37
+ prediction = self.multi_model[model_id](inputs, **parameters)
38
+ else:
39
+ prediction = self.multi_model[model_id](inputs)
40
+ # postprocess the prediction
41
+ return prediction