Update pipeline.py
Browse files- 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,
|
81 |
# process input
|
82 |
-
inputs =
|
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}
|