change infer function

infer_concat.py

@@ -63,7 +63,7 @@ def processing_data_infer(input_file):
 tokenizer = AutoTokenizer.from_pretrained("VietAI/vit5-base-vietnews-summarization")
 model = AutoModelForSeq2SeqLM.from_pretrained("VietAI/vit5-base-vietnews-summarization")
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
 model.to(device)
 
 model.load_state_dict(torch.load("./weight_cp19_model.pth", map_location=torch.device('cpu')))
@@ -90,12 +90,12 @@ def infer_2_hier(model, data_loader, device, tokenizer):
                 summary = model.generate(inputs[i].to(device),
                                          attention_mask=att_mask[i].to(device),
                                          max_length=128,
-                                         num_beams=12,
-                                         num_return_sequences=1)
+                                         num_beams=4,
+                                         num_return_sequences=1, no_repeat_ngram_size=3)
                 summaries.append(summary)
         summaries = torch.cat(summaries, dim = 1)
-        for k in summaries:
-                all_summaries.append(tokenizer.decode(k, skip_special_tokens=True))
+	
+        all_summaries.append(tokenizer.decode(summaries, skip_special_tokens=True))
 
     
     end = time.time()
@@ -104,6 +104,6 @@ def infer_2_hier(model, data_loader, device, tokenizer):
 
 def vit5_infer(data):
 	dataset = Dataset4Summarization(data, tokenizer)
-	data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=1)
+	data_loader = torch.utils.data.DataLoader(dataset, batch_size=1)
 	result = infer_2_hier(model, data_loader, device, tokenizer)
 	return result