Lin0He commited on
Commit
0d18729
·
1 Parent(s): 343d546

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +2 -2
pipeline.py CHANGED
@@ -77,9 +77,9 @@ class PreTrainedPipeline():
77
  self.model = AutoModelForCausalLM.from_pretrained("Lin0He/text-summary-gpt2-short")
78
 
79
 
80
- def __call__(self, data) -> Dict[str, Any]:
81
  # process input
82
- inputs = "Today was a beautiful day. I had a good night’s sleep and was ready to start the day. I went to work and had a productive morning. I even managed to finish a project I’d been working on for weeks. After work, I ran to clear my head. It was a beautiful day and the weather was perfect for it. I came home and cooked dinner with my partner. We had a nice conversation over dinner and then spent the evening watching a movie. Overall, it was a pretty relaxing and enjoyable day." #data.pop("inputs", data)
83
  # process input text
84
  prediction = model_infer(self.model, self.tokenizer, inputs+"TL;DR")
85
  return {"text": prediction}
 
77
  self.model = AutoModelForCausalLM.from_pretrained("Lin0He/text-summary-gpt2-short")
78
 
79
 
80
+ def __call__(self, Dict[str, str]) -> Dict[str, str]:
81
  # process input
82
+ inputs = data.pop("inputs", data)
83
  # process input text
84
  prediction = model_infer(self.model, self.tokenizer, inputs+"TL;DR")
85
  return {"text": prediction}