philschmid HF staff commited on
Commit
ae6b49a
1 Parent(s): b3f4606

Delete handler.py

Browse files
Files changed (1) hide show
  1. handler.py +0 -33
handler.py DELETED
@@ -1,33 +0,0 @@
1
- from typing import Dict, List, Any
2
- from transformers import pipeline, AutoTokenizer
3
- import torch
4
- import os
5
-
6
- os.system("python -m torch.utils.collect_env")
7
-
8
- class EndpointHandler():
9
- def __init__(self, path=""):
10
- # load the optimized model
11
- self.pipeline = pipeline("text-classification", model=path)
12
-
13
-
14
- def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
15
- """
16
- Args:
17
- data (:obj:):
18
- includes the input data and the parameters for the inference.
19
- Return:
20
- A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
21
- - "label": A string representing what the label/class is. There can be multiple labels.
22
- - "score": A score between 0 and 1 describing how confident the model is for this label/class.
23
- """
24
- inputs = data.pop("inputs", data)
25
- parameters = data.pop("parameters", None)
26
-
27
- # pass inputs with all kwargs in data
28
- if parameters is not None:
29
- prediction = self.pipeline(inputs, **parameters)
30
- else:
31
- prediction = self.pipeline(inputs)
32
- # postprocess the prediction
33
- return prediction