Lin0He commited on
Commit
6d55152
·
1 Parent(s): e055f73

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +2 -2
pipeline.py CHANGED
@@ -65,9 +65,9 @@ class PreTrainedPipeline():
65
  self.tokenizer = AutoTokenizer.from_pretrained("Lin0He/text-summary-gpt2-short")
66
  self.model = AutoModelForCausalLM.from_pretrained("Lin0He/text-summary-gpt2-short")
67
 
68
- def __call__(self, data) -> Dict[str, Any]:
69
  # process input
70
  inputs = data#.pop("inputs", data)
71
  # process input text
72
  prediction = model_infer(self.model, self.tokenizer, inputs+"TL;DR")
73
- return {"text": prediction}
 
65
  self.tokenizer = AutoTokenizer.from_pretrained("Lin0He/text-summary-gpt2-short")
66
  self.model = AutoModelForCausalLM.from_pretrained("Lin0He/text-summary-gpt2-short")
67
 
68
+ def __call__(self, data) -> str:
69
  # process input
70
  inputs = data#.pop("inputs", data)
71
  # process input text
72
  prediction = model_infer(self.model, self.tokenizer, inputs+"TL;DR")
73
+ return prediction