ydshieh
/
vit-gpt2-coco-en
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update pipeline.py

This commit is contained in:
Yih-Dar SHIEH 2021-10-25 08:30:24 +00:00 committed by huggingface-web
parent 45bde45ae6
commit 462822078f
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
pipeline.py
View File

@ -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