Update pipeline.py
Browse files- pipeline.py +9 -5
pipeline.py
CHANGED
@@ -20,7 +20,7 @@ class PreTrainedPipeline():
|
|
20 |
max_length = 16
|
21 |
num_beams = 4
|
22 |
# self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
|
23 |
-
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True}
|
24 |
|
25 |
self.model.to("cpu")
|
26 |
self.model.eval()
|
@@ -29,8 +29,12 @@ class PreTrainedPipeline():
|
|
29 |
def _generate(pixel_values):
|
30 |
|
31 |
with torch.no_grad():
|
32 |
-
|
33 |
-
|
|
|
|
|
|
|
|
|
34 |
|
35 |
self.generate = _generate
|
36 |
|
@@ -49,10 +53,10 @@ class PreTrainedPipeline():
|
|
49 |
# pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values
|
50 |
pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values
|
51 |
|
52 |
-
output_ids = self.generate(pixel_values)
|
53 |
preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
54 |
preds = [pred.strip() for pred in preds]
|
55 |
|
56 |
-
preds = [{"label": preds[0], "score":
|
57 |
|
58 |
return preds
|
|
|
20 |
max_length = 16
|
21 |
num_beams = 4
|
22 |
# self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
|
23 |
+
self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True, "output_scores": True}
|
24 |
|
25 |
self.model.to("cpu")
|
26 |
self.model.eval()
|
|
|
29 |
def _generate(pixel_values):
|
30 |
|
31 |
with torch.no_grad():
|
32 |
+
|
33 |
+
outputs = self.model.generate(pixel_values, **self.gen_kwargs)
|
34 |
+
output_ids = outputs.sequences
|
35 |
+
sequences_scores = outputs.sequences_scores
|
36 |
+
|
37 |
+
return output_ids, sequences_scores
|
38 |
|
39 |
self.generate = _generate
|
40 |
|
|
|
53 |
# pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values
|
54 |
pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values
|
55 |
|
56 |
+
output_ids, sequences_scores = self.generate(pixel_values)
|
57 |
preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
58 |
preds = [pred.strip() for pred in preds]
|
59 |
|
60 |
+
preds = [{"label": preds[0], "score": float(sequences_scores[0])}]
|
61 |
|
62 |
return preds
|