File size: 14,601 Bytes
e56055d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
441
442
import datetime
import typing
import numpy as np
import struct
import os
import getpass
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 = '/iris/u/clin/code/efk/'
        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, minimize: bool = False):
        self.best_value = 1e9 if minimize else -1e9
        self.best_iter = 0
        self.current_iter = 0
        self.key = key
        self.patience = patience
        self.minimize = minimize
        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 self.minimize else 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).mul(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()