hungdungn47 commited on
Commit
d088bc3
·
1 Parent(s): 48742d4
Files changed (1) hide show
  1. infer_concat.py +3 -1
infer_concat.py CHANGED
@@ -100,10 +100,12 @@ def infer_concat(model, data_loader, device, tokenizer):
100
 
101
  end = time.time()
102
  print(f"Time: {end-start}")
 
103
  return all_summaries
104
 
105
  def vit5_infer(data):
 
106
  dataset = Dataset4Summarization(data, tokenizer)
107
- data_loader = torch.utils.data.DataLoader(dataset, batch_size=1)
108
  result = infer_concat(model, data_loader, device, tokenizer)
109
  return result
 
100
 
101
  end = time.time()
102
  print(f"Time: {end-start}")
103
+ print(all_summaries)
104
  return all_summaries
105
 
106
  def vit5_infer(data):
107
+ print(data)
108
  dataset = Dataset4Summarization(data, tokenizer)
109
+ data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=4)
110
  result = infer_concat(model, data_loader, device, tokenizer)
111
  return result