sort test_file; remove labels in the collate_fn
@@ -97,7 +97,7 @@ class ExtSumProcessedDataset(Dataset):
files is shuffled when the dataset is loaded. Defaults to False.
"""
- self.file_list = file_list
+ self.file_list = sorted(file_list)
if is_shuffle:
random.shuffle(file_list)
self.data = []
@@ -650,22 +650,10 @@ class ExtractiveSummarizer(Transformer):
# tuple_batch = [list(col) for col in zip(*[d.values() for d in dict_list]
if dict_list is None or len(dict_list) <= 0:
return None
- is_labeled = False
- if "labels" in dict_list[0]:
- is_labeled = True
tuple_batch = [list(d.values()) for d in dict_list]
## generate mask and mask_cls, and only select tensors for the model input
- batch = Batch(tuple_batch, is_labeled=True)
- if is_labeled:
- return {
- "src": batch.src,
- "segs": batch.segs,
- "clss": batch.clss,
- "mask": batch.mask,
- "mask_cls": batch.mask_cls,
- "labels": batch.labels,
- }
- else:
+ ## the labels was never used in prediction, set is_labeled as False
+ batch = Batch(tuple_batch, is_labeled=False)
return {
"src": batch.src,
"segs": batch.segs,
