Update pipeline.py
Browse files- 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) ->
|
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
|
|
|
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
|