Update pipeline.py

2021-10-25 08:19:08 +00:00 · 2021-10-25 08:19:08 +00:00 · 45bde45ae6
parent 117d46a115
commit 45bde45ae6
1 changed files with 2 additions and 0 deletions
--- a/pipeline.py
+++ b/pipeline.py
@ -53,4 +53,6 @@ class PreTrainedPipeline():
        preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
        preds = [pred.strip() for pred in preds]

+        preds = [{"label": preds[0], "score": 1.0}]
+
        return preds