-
Notifications
You must be signed in to change notification settings - Fork 6
/
utils.py
441 lines (342 loc) · 14.2 KB
/
utils.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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
import datetime
import typing
import numpy as np
import struct
import os
import getpass
import hydra
import logging
import torch
import torch.nn as nn
from collections import defaultdict
import math
LOG = logging.getLogger(__name__)
def masked_mean(values, mask):
assert mask.dtype == torch.bool
assert values.shape == mask.shape
return (values * mask.float()).sum() / mask.sum().float()
def mask_hf_labels(labels, null_token=0):
valid_mask = labels != -100
valid_labels = labels.masked_fill(~valid_mask, null_token)
return valid_mask, valid_labels
def gather_log_probs(logits, labels):
assert labels.dim() == logits.dim() - 1
assert labels.shape == logits.shape[:-1]
return logits.log_softmax(-1).gather(-1, labels.unsqueeze(-1)).squeeze(-1)
def off_diagonal(mat):
assert mat.dim() == 2
# assert mat.shape[0] == mat.shape[1]
mask = ~torch.eye(max(mat.shape), dtype=torch.bool)
mask = mask[:mat.shape[0], :mat.shape[1]]
off_d = mat[mask]
assert off_d.numel() == mat.shape[0] * mat.shape[1] - min(mat.shape)
return off_d
def set_dropout(model, p):
if p is not None:
n_reset = 0
for m in model.modules():
if isinstance(m, nn.Dropout):
m.p = p
n_reset += 1
if hasattr(m, "dropout"): # Requires for BART, which uses F.dropout
if isinstance(m.dropout, float):
m.dropout = p
n_reset += 1
if hasattr(m, "activation_dropout"): # Requires for BART, which uses F.dropout
if isinstance(m.activation_dropout, float):
m.activation_dropout = p
n_reset += 1
LOG.info(f"Set {n_reset} dropout modules to p={p}")
def _inner_params(named_parameters, inner_names):
param_dict = dict(named_parameters)
return [(n, param_dict[n]) for n in inner_names]
def shift_targets(config):
return "t5" not in config.model.name.lower() and "blender" not in config.model.name.lower()
# https://stackoverflow.com/questions/32871539/integer-factorization-in-python
def factorization(n):
return [(i, n // i) for i in range(1, int(n**0.5) + 1) if n % i == 0]
def scr():
if os.path.exists("/scr-ssd"):
scr_dir = "/scr-ssd/" + getpass.getuser()
else:
scr_dir = "/scr/" + getpass.getuser()
if not os.path.exists(scr_dir):
os.makedirs(scr_dir)
return scr_dir
def uuid(digits=4):
if not hasattr(uuid, "uuid_value"):
uuid.uuid_value = struct.unpack('I', os.urandom(4))[0] % int(10**digits)
return uuid.uuid_value
def formatted_timestamp(time=None):
if time is None:
time = datetime.datetime.now()
return time.strftime("%d/%m/%Y-%H:%M:%S/%f")
def time_delta_seconds(start, finish=None):
assert type(start) == str
t1 = datetime.datetime.strptime(start, "%d/%m/%Y-%H:%M:%S/%f")
if finish is not None:
assert type(finish) == str
t2 = datetime.datetime.strptime(finish, "%d/%m/%Y-%H:%M:%S/%f")
else:
t2 = datetime.datetime.now()
return (t2 - t1).total_seconds()
def dict_to(d, device):
new_dict = {}
for k, v in d.items():
if isinstance(v, torch.Tensor):
new_dict[k] = v.to(device)
elif isinstance(v, dict):
new_dict[k] = dict_to(v, device)
else:
new_dict[k] = v
return new_dict
def safe_backward(loss, parameters, accumulate=1, allow_unused=False, backward=False):
if backward:
(loss / accumulate).backward()
else:
parameters = list(parameters) # Capture the generator output
grads = torch.autograd.grad(loss, parameters, allow_unused=allow_unused)
nan, inf = False, False
for g in grads:
if g is not None:
nan |= g.isnan().any().item()
inf |= g.isinf().any().item()
if not (nan or inf):
for p, g in zip(parameters, grads):
if g is None:
continue
if p.grad is None:
p.grad = g / accumulate
else:
p.grad += g / accumulate
else:
LOG.info(f"Skipping grad accumulation because inf: {inf} nan: {nan}")
def _logits(x):
return x if not hasattr(x, "logits") else x.logits
def _last_encoder_state(x):
if hasattr(x, "encoder_last_hidden_state"):
return x.encoder_last_hidden_state
else:
return x.hidden_states[-1]
def load_archive(path):
import torch
if not os.path.exists(path):
# We've not passed an explicit path, but a part of the filename
wd = hydra.utils.get_original_cwd()
directories = ["outputs", "multirun"]
matches = []
for d in directories:
search = os.path.join(wd, d)
for run_dir in os.listdir(search):
if path in run_dir:
matches.append(os.path.join(search, run_dir))
assert len(matches) == 1, f">1 matches for search {path}; specify exact path"
full_run_dir = matches[0]
if "0" in os.listdir(full_run_dir):
full_run_dir = os.path.join(full_run_dir, "0")
models_dir = os.path.join(full_run_dir, "models")
models = os.listdir(models_dir)
non_bk = [m for m in models if not m.endswith(".bk")]
assert (
len(non_bk) == 1
), f"Expected a single model in {models_dir}, got {len(non_bk)}"
path = os.path.join(models_dir, non_bk[0])
LOG.info(f"Loading checkpoint from {path}")
archive = torch.load(path, map_location="cpu")
LOG.info("Load complete.")
return archive, path
def flatten_dict(d):
to_process = list(d.items())
output = {}
while len(to_process):
k, v = to_process.pop()
if isinstance(v, typing.MutableMapping):
to_process.extend([(f"{k}.{k_}", v_) for (k_, v_) in v.items()])
else:
assert k not in output.keys(), "Somehow ended up with duplicate keys"
output[k] = v
return output
def add_padding(tokenizer, model):
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))
model.transformer.wte.weight.data[-1] = model.transformer.wte.weight.data.mean(0)
def add_sep(tokenizer, model):
tokenizer.add_special_tokens({'sep_token': '[SEP]'})
# model.resize_token_embeddings(len(tokenizer))
# model.lm_head.weight.data[-1, :] = model.lm_head.weight.data.mean(0)
class EarlyStopper:
def __init__(self, patience: int, key: str):
self.best_value = 1e9
self.best_iter = 0
self.current_iter = 0
self.key = key
self.patience = patience
self._stop = False
def update(self, idx, stats):
assert self.key in stats, f"'{self.key}' not in stats dict"
value = stats[self.key]
new_best = value < self.best_value
if new_best:
self.best_value = value
self.best_iter = idx
self.current_iter = idx
return new_best
def should_stop(self):
self._stop |= self.current_iter - self.best_iter >= self.patience
return self._stop
class RunningStatAverager:
def __init__(self, suffix="", exclude=["grad/"], compute_ppl: bool = True):
self.underlying = None
self.suffix = suffix
self.exclude = exclude
self.compute_ppl = compute_ppl
self.reset()
def add(self, d: dict):
for k, v in d.items():
if not any([k.startswith(prefix) for prefix in self.exclude]):
if len(self.suffix):
self.underlying[f"{k}_{self.suffix}"].append(v)
else:
self.underlying[k].append(v)
def average(self):
average = {}
for k, v in self.underlying.items():
if not k.startswith("nll/"):
average[k] = sum(v) / len(v)
else:
assert len(k.split("/")) == 2, f"Invalid key {k}"
name = k.split("/")[1]
token_counts = self.underlying[f"n_tokens/{name}"]
total_nll = sum([nll * c for nll, c in zip(v, token_counts)])
average[k] = total_nll / sum(token_counts)
if self.compute_ppl:
average[f"perplexity/{name}"] = math.e ** average[k]
return {k: v if not isinstance(v, torch.Tensor) else v.item() for k, v in average.items()}
def reset(self):
self.underlying = defaultdict(list)
class EditBatchSampler:
def __init__(
self,
n,
memorize_mode=False,
loc_disjoint=True,
seed=0,
hard_neg=False,
hard_neg_prob=1.0,
loc_distr_matrix=None,
loc_idx_matrix=None,
keep_probs=None,
mutex=None
):
self.memorize_mode = memorize_mode
self.n = n
self.loc_disjoint = loc_disjoint
self.rng = np.random.default_rng(seed)
self.hard_neg = hard_neg
self.hard_neg_prob = hard_neg_prob
self.loc_probs = loc_distr_matrix
self.loc_idxs = loc_idx_matrix
self.keep_probs = np.array(keep_probs)[:self.n] if keep_probs is not None else None
self.mutex = mutex[:self.n] if mutex is not None else None
self._init()
def _init(self):
idxs = np.arange(self.n)
if self.keep_probs is not None:
sample = self.rng.binomial(1, self.keep_probs).astype(np.bool)
idxs = idxs[sample]
self.perm = self.rng.permutation(idxs)
self.edit_position = 0
def get_edit_idxs(self, batch_size):
if self.mutex is None:
idxs = set([int(idx) for idx in self.perm[self.edit_position: self.edit_position + batch_size]])
self.edit_position += batch_size
else:
mutexes = []
idxs = []
def notin(x, mutexes):
for m in mutexes:
if x in m or m in x:
return False
return True
while len(idxs) < batch_size:
new_idx = self.perm[self.edit_position]
if notin(self.mutex[new_idx], mutexes):
mutexes.append(self.mutex[new_idx])
idxs.append(int(new_idx))
self.edit_position += 1
if self.edit_position == self.perm.shape[0]:
return None
idxs = set(idxs)
return idxs
def sample(self, batch_size, return_hard_flag=False):
if self.memorize_mode:
return list(range(batch_size)), list(range(batch_size, batch_size * 2))
if self.edit_position + batch_size >= self.perm.shape[0]:
self._init() # Re-start if we end with a partially-sized batch
edit_idxs = self.get_edit_idxs(batch_size)
if edit_idxs is None:
self._init()
edit_idxs = self.get_edit_idxs(batch_size)
if edit_idxs is None:
raise RuntimeError(f"No valid batches of size {batch_size} exist!")
if self.hard_neg:
assert self.loc_probs is not None, "hard_neg is on, but don't have distance matrix!"
def get_loc_idxs():
if self.hard_neg and self.rng.uniform() < self.hard_neg_prob:
return [int(self.rng.choice(self.loc_idxs[idx], p=self.loc_probs[idx])) for idx in edit_idxs], True
else:
# Use deterministic implementation in case edit batches are large
non_edit_idxs = list(set(range(self.n)) - set(edit_idxs))
return [int(idx) for idx in self.rng.choice(non_edit_idxs, batch_size)], False
loc_idxs, hard = get_loc_idxs()
if self.loc_disjoint:
steps = 0
while len(edit_idxs.intersection(set(loc_idxs))) > 0:
loc_idxs, hard = get_loc_idxs()
steps += 1
if steps > 100:
raise RuntimeError("Can't find disjoint loc_idxs and edit_idxs!")
if return_hard_flag:
return list(edit_idxs), loc_idxs, hard
else:
return list(edit_idxs), loc_idxs
def parent_module(model, pname):
comps = pname.split('.')
parent = model
for comp in comps[:-1]:
if hasattr(parent, comp):
parent = getattr(parent, comp)
elif comp.isdigit():
parent = parent[int(comp)]
else:
raise RuntimeError(f"Couldn't find child module {comp}")
assert hasattr(parent, comps[-1])
return parent
def build_distr_matrix(edit_qs, config, loc_qs=None, slice_size=1000):
n = len(edit_qs)
device = "cuda" if torch.cuda.is_available() else "cpu"
num_neighbors = config.data.hard_neg_neighbors
num_exclude = config.data.hard_neg_exclude
temp = config.data.hard_neg_temp
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import pytorch_cos_sim
embedding_model = SentenceTransformer('all-MiniLM-L6-v2', cache_folder=scr()).to(device)
ind_matrix = torch.zeros((n, num_neighbors - num_exclude), dtype=torch.long)
distr_matrix = torch.full((n, num_neighbors - num_exclude), float('nan'))
edit_encodings = torch.FloatTensor(embedding_model.encode(edit_qs, batch_size=256)).to(device)
# If loc_qs is None then build the similarity matrix between edit_qs and itself
loc_encodings = edit_encodings if loc_qs is None else embedding_model.encode(loc_qs, batch_size=256)
if isinstance(loc_encodings, np.ndarray):
loc_encodings = torch.FloatTensor(loc_encodings).to(device)
for idx in range(0, n, slice_size):
end_idx = idx + slice_size if idx + slice_size <= n else n
slice_encodings = edit_encodings[idx:end_idx]
sim_rows = pytorch_cos_sim(slice_encodings, loc_encodings)
indices = sim_rows.topk(num_neighbors, -1).indices[:, num_exclude:]
ind_matrix[idx:end_idx] = indices.cpu()
distr_matrix[idx:end_idx] = sim_rows.gather(-1, indices).div(temp).exp().cpu()
assert not torch.isnan(distr_matrix).any()
LOG.info(f"Built hard negative distribution matrix of size {distr_matrix.shape}")
distr_matrix = distr_matrix.numpy()
distr_matrix = distr_matrix / distr_matrix.sum(-1, keepdims=True)
return distr_matrix, ind_matrix.numpy()