-
Notifications
You must be signed in to change notification settings - Fork 1
/
collate.py
executable file
·94 lines (74 loc) · 4.13 KB
/
collate.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
import logging
import torch
import math
from util import pad_clusters
logger = logging.getLogger(__name__)
class SegmentCollator:
def __init__(self, tokenizer, device, max_segment_len):
self.tokenizer = tokenizer
self.device = device
self.max_segment_len = max_segment_len
def __call__(self, batch):
# pad to the longest doc in the batch
batch = self.tokenizer.pad(batch)
batch['leftovers'] = {'input_ids': [], 'attention_mask': []}
# break down to segment of segment len
input_ids = [[ids[i:i + self.max_segment_len] for i in range(0, len(ids), self.max_segment_len)]
for ids in batch['input_ids']]
attention_mask = [[mask[i:i + self.max_segment_len] for i in range(0, len(mask), self.max_segment_len)]
for mask in batch['attention_mask']]
# if we have more than 1 segment and the last segment is less than segment_len we have leftovers.
if len(input_ids[0]) > 1 and len(input_ids[0][-1]) < self.max_segment_len:
batch['leftovers']['input_ids'] = torch.tensor([ids[-1] for ids in input_ids], device=self.device)
batch['leftovers']['attention_mask'] = torch.tensor([mask[-1] for mask in attention_mask], device=self.device)
# remove leftovers from main batch
input_ids = [ids[:-1] for ids in input_ids]
attention_mask = [mask[:-1] for mask in attention_mask]
batch['input_ids'] = torch.tensor(input_ids, device=self.device)
batch['attention_mask'] = torch.tensor(attention_mask, device=self.device)
# if all the docs with no clusters pad it to 1 cluster with 1 mention.
max_num_clusters, max_max_cluster_size = max(1, max(batch['num_clusters'])), max(1, max(batch['max_cluster_size']))
padded_clusters = [pad_clusters(cluster, max_num_clusters, max_max_cluster_size) for cluster in batch['gold_clusters']]
batch['gold_clusters'] = torch.tensor(padded_clusters, device=self.device)
return batch
class LongformerCollator:
def __init__(self, tokenizer, device):
self.tokenizer = tokenizer
self.device = device
def __call__(self, batch):
# pad to the longest doc in the batch
batch = self.tokenizer.pad(batch)
batch['input_ids'] = torch.tensor(batch['input_ids'], device=self.device)
batch['attention_mask'] = torch.tensor(batch['attention_mask'], device=self.device)
max_num_clusters, max_max_cluster_size = max(batch['num_clusters']), max(batch['max_cluster_size'])
if max_num_clusters == 0:
batch['gold_clusters'] = None
else:
padded_clusters = [pad_clusters(cluster, max_num_clusters, max_max_cluster_size) for cluster in batch['gold_clusters']]
batch['gold_clusters'] = torch.tensor(padded_clusters, device=self.device)
return batch
class DynamicBatchSampler:
def __init__(self, dataset, collator, max_tokens, max_segment_len, max_doc_len=None):
self.max_tokens = max_tokens
self.dataset = dataset.sort('length', reverse=True)
self.collator = collator
self.max_segment_len = max_segment_len
self.max_doc_len = max_doc_len
def __iter__(self):
batch = []
per_example_batch_len = 0
for example in self.dataset:
if self.max_doc_len is not None and example['length'] > self.max_doc_len:
logger.info(f'Skipping doc with len {example["length"]}. max_doc_len is {self.max_doc_len}')
continue
if not batch:
per_example_batch_len = self.calc_effective_per_example_batch_len(example['length'])
elif (len(batch) + 1) * per_example_batch_len > self.max_tokens:
yield self.collator(batch)
batch = []
per_example_batch_len = self.calc_effective_per_example_batch_len(example['length'])
batch.append(example)
if len(batch) > 0:
yield self.collator(batch)
def calc_effective_per_example_batch_len(self, example_len):
return math.ceil(example_len / self.max_segment_len) * self.max_segment_len