Update pipeline.py

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