From 462822078f34dfe49d8ed7061444a6ab327b8ca5 Mon Sep 17 00:00:00 2001 From: Yih-Dar SHIEH Date: Mon, 25 Oct 2021 08:30:24 +0000 Subject: [PATCH] Update pipeline.py --- pipeline.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/pipeline.py b/pipeline.py index 1d9894b..41f625f 100644 --- a/pipeline.py +++ b/pipeline.py @@ -20,7 +20,7 @@ class PreTrainedPipeline(): max_length = 16 num_beams = 4 # self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams} - self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True} + self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True, "output_scores": True} self.model.to("cpu") self.model.eval() @@ -29,8 +29,12 @@ class PreTrainedPipeline(): def _generate(pixel_values): with torch.no_grad(): - output_ids = self.model.generate(pixel_values, **self.gen_kwargs).sequences - return output_ids + + outputs = self.model.generate(pixel_values, **self.gen_kwargs) + output_ids = outputs.sequences + sequences_scores = outputs.sequences_scores + + return output_ids, sequences_scores self.generate = _generate @@ -49,10 +53,10 @@ class PreTrainedPipeline(): # pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values - output_ids = self.generate(pixel_values) + output_ids, sequences_scores = self.generate(pixel_values) 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}] + preds = [{"label": preds[0], "score": float(sequences_scores[0])}] return preds