philschmid HF staff commited on
Commit
4762d81
1 Parent(s): e2742b6

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +3 -11
pipeline.py CHANGED
@@ -12,10 +12,10 @@ class PreTrainedPipeline():
12
  self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
13
 
14
 
15
- def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
16
  """
17
  Args:
18
- data (:obj:`str`):
19
  a string containing some text
20
  Return:
21
  A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
@@ -23,12 +23,4 @@ class PreTrainedPipeline():
23
  - "score": A score between 0 and 1 describing how confident the model is for this label/class.
24
  """
25
  # pop inputs for pipeline
26
- inputs = data.pop("inputs", data)
27
- parameters = data.pop("parameters", None)
28
-
29
- # pass inputs with all kwargs in data
30
- if parameters is not None:
31
- prediction = self.pipeline(inputs, **parameters)
32
- else:
33
- prediction = self.pipeline(inputs)
34
- return prediction
 
12
  self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
13
 
14
 
15
+ def __call__(self, inputs: Any) -> List[List[Dict[str, float]]]:
16
  """
17
  Args:
18
+ inputs (:obj:`str`):
19
  a string containing some text
20
  Return:
21
  A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
 
23
  - "score": A score between 0 and 1 describing how confident the model is for this label/class.
24
  """
25
  # pop inputs for pipeline
26
+ return self.pipeline(inputs)