kwang2049 commited on
Commit
5a19953
1 Parent(s): 5e0e65c

input type

Browse files
Files changed (1) hide show
  1. pipeline.py +2 -4
pipeline.py CHANGED
@@ -17,7 +17,7 @@ class PreTrainedPipeline:
17
  self.predictor = CorefPredictor.from_extracted_archive(archive_content)
18
  self.predictor.set_device("cpu")
19
 
20
- def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
21
  """
22
  data args:
23
  inputs (:obj: `str`)
@@ -26,7 +26,5 @@ class PreTrainedPipeline:
26
  A :obj:`list` | `dict`: will be serialized and returned
27
  """
28
  # get inputs
29
- return data
30
- inputs: str = data.pop("inputs", data)
31
- prediction = self.predictor.resolve_paragraphs(inputs.split("\n\n"))
32
  return prediction.to_dict()
 
17
  self.predictor = CorefPredictor.from_extracted_archive(archive_content)
18
  self.predictor.set_device("cpu")
19
 
20
+ def __call__(self, data: str) -> Dict[str, Any]:
21
  """
22
  data args:
23
  inputs (:obj: `str`)
 
26
  A :obj:`list` | `dict`: will be serialized and returned
27
  """
28
  # get inputs
29
+ prediction = self.predictor.resolve_paragraphs(data.split("\n\n"))
 
 
30
  return prediction.to_dict()