| import transformers |
| import torch |
| import os |
| import struct |
| import random |
|
|
| CONTEXT_TEMPLATES_CACHE = None |
|
|
| def find_sublist_start_index(list1, list2): |
| for i in range(len(list1) - len(list2)+1): |
| if all(a == b for a, b in zip(list1[i:i+len(list2)], list2)): |
| return i |
| return None |
|
|
| def get_inner_params(named_parameters, inner_names): |
| param_dict = dict(named_parameters) |
| return [(n, param_dict[n]) for n in inner_names] |
|
|
| def param_subset(named_parameters, inner_names): |
| param_dict = dict(named_parameters) |
| return [param_dict[n] for n in inner_names] |
|
|
| def print_trainable_parameters(model, new_weight, mask_ratio): |
| original_parameters = 0 |
| new_weight_param = 0 |
| for _, param in new_weight.named_parameters(): |
| new_weight_param += param.numel() |
| for _, param in model.named_parameters(): |
| original_parameters += param.numel() |
| print(f"Original Model params: {original_parameters} || New Weight params: {new_weight_param} || trainable%: {100 * new_weight_param * (1-mask_ratio) / original_parameters}") |
|
|
|
|
| def parent_module(model, pname): |
| components = pname.split('.') |
| parent = model |
|
|
| for component in components[:-1]: |
| if hasattr(parent, component): |
| parent = getattr(parent, component) |
| elif component.isdigit(): |
| parent = parent[int(component)] |
| else: |
| raise RuntimeError(f"Couldn't find child module {component}") |
|
|
| if not hasattr(parent, components[-1]): |
| raise RuntimeError(f"Couldn't find child module {components[-1]}") |
|
|
| return parent |
|
|
| 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 ckpt_dir(): |
| """returns the directory in which to store model checkpoints""" |
| path = "./ckpts/" |
| if not os.path.exists(path): |
| os.makedirs(path) |
| return path |
|
|
| def brackets_to_periods(name): |
| return name.replace("[", ".").replace("]", "") |
| |
| def get_params(model): |
| return model.state_dict() |
|
|
| def get_shape(p, model): |
| |
| return p.shape if isinstance(model, transformers.GPT2LMHeadModel) else (p.shape[1], p.shape[0]) |
|
|
| def get_logits(x): |
| return x.logits if hasattr(x, "logits") else x |
|
|
|
|
| LOC_PROMPTS = ['nq question: who played mr grainger in are you being served Arthur Brough', |
| "nq question: who sings the song let's hear it for the boy Deniece Williams", |
| "nq question: who wrote all my ex's live in texas Sanger D. Shafer", |
| "nq question: when is the america's got talent finale 2018 September 19, 2018", |
| "nq question: what is the fifth biggest state in the united states New Mexico", |
| "nq question: who plays john black on days of our lives Drake Hogestyn (/ˈhʌdʒstən/; born Donald Drake Hogestyn", |
| "nq question: what is the name of the new star wars movie The Last Jedi", |
| "nq question: what is the main principle of path-goal theory a leader's behavior is contingent to the satisfaction, motivation and performance of his or her subordinates", |
| "nq question: who plays luna's dad in harry potter Ifans", |
| "nq question: who has the most grammy nominations as an artist Quincy Jones", |
| "nq question: what is the control unit function in the cpu tells the computer's memory, arithmetic/logic unit and input and output devices how to respond to the instructions that have been sent to the processor", |
| "nq question: who was the first indian prime minister to visit palestine Narendra Modi", |
| "nq question: where did the plane carrying the marshall football team crash into a hill just short of the Tri-State Airport", |
| "nq question: what movie is the line lighten up francis from Stripes", |
| "nq question: set of rules for solving a mathematical or computational problem in finite number of steps an algorithm", |
| "nq question: who changed indian capital from calcutta to delhi George V", |
| "nq question: who did bette midler play in the rose Mary Rose Foster (The Rose)", |
| "nq question: how much did it cost to make the new star wars movie $200–217 million" |
| ] |
|
|
| def tokenize(batch, tokenizer, device, context_templates=None, hparams=None): |
| prompt, label = batch["prompt"], batch["target_new"] |
| batch['loc_prompt'] = random.choice(LOC_PROMPTS) |
| if not isinstance(prompt, list): |
| prompt=[prompt] |
| if not isinstance(label, list): |
| label=[label] |
| mask_token = -100 |
|
|
| |
| full_prompt = [f"{templ.format(p + ' ' + l)}" for p, l in zip(prompt, label) for templ in context_templates] |
| full_prompt += [batch['loc_prompt']] |
|
|
| prompt_ids = tokenizer([f"{templ.format(p)}" for p in prompt for templ in context_templates], return_tensors="pt", padding=True, truncation=True)["input_ids"] |
|
|
| num_prompt_toks = [len(i) for i in prompt_ids] |
| tokens = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True) |
| tokens["labels"] = tokens["input_ids"].clone() |
| if hparams.objective_optimization == 'only_label': |
| for i in range(len(num_prompt_toks)): |
| tokens["labels"][i][:num_prompt_toks[i]] = mask_token |
|
|
| tokens["labels"][tokens["input_ids"] == tokenizer.pad_token_id] = mask_token |
| if batch['loc_prompt'] in batch['prompt']: |
| subject_token = tokenizer.encode(' ' + batch['loc_prompt'], add_special_tokens=False) |
| subject_token1 = tokenizer.encode(batch['loc_prompt'], add_special_tokens=False) |
| subject_length = len(subject_token) |
| act_mask = torch.zeros_like(tokens['input_ids'][:-1]) |
| deact_mask = torch.zeros_like(tokens['input_ids'][:-1]) |
| for i, token in enumerate(tokens['input_ids'][:-1]): |
| start_idx = find_sublist_start_index(token.detach().cpu().numpy().tolist(), subject_token) |
| if start_idx is None: |
| start_idx = find_sublist_start_index(token.detach().cpu().numpy().tolist(), subject_token1) |
| subject_length = len(subject_token1) |
| act_mask[i][start_idx: start_idx + subject_length] = 1 |
| deact_mask[i][:start_idx] = 1 |
| deact_mask[i][start_idx + subject_length:] = 1 |
|
|
| act_mask = act_mask.to(device) |
| deact_mask = deact_mask.to(device) |
| else: |
| act_mask = None |
| deact_mask = None |
|
|
| tokens = {f"{k1}" : v1.to(device) for k1, v1 in tokens.items()} |
| return tokens, act_mask, deact_mask |
|
|
| class EarlyStopMeter: |
| """Computes and stores the average and current value""" |
|
|
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.avg = 0 |
| self.pre = 0 |
| self.val = 1e9 |
| self.sum = 0 |
| self.count = 0 |
|
|
| def update(self, val): |
| self.pre = self.val |
| self.val = val |
| self.sum += val |
| self.count += 1 |
| self.avg = self.sum / self.count |
|
|
| def stop(self, ): |
| return abs(self.val - self.pre) <= 1e-4 and self.val <= 0.02 |
|
|
| class EditingMeanAct: |
| """Computes and stores the average and current value""" |
|
|
| def __init__(self, min_a=1e9): |
| self.reset(min_a=min_a) |
|
|
| def reset(self, min_a=1e9): |
| self.avg = 0 |
| self.count = 0 |
| self.sum = 0 |
| self.min_a = min_a |
|
|
| def update(self, val): |
| self.sum += val |
| self.count += 1 |
| self.avg = self.sum / self.count |
| self.min_a = min(self.min_a, val) |
|
|
| def mean_act(self): |
| return self.avg |
| def min_act(self): |
| return self.min_a |
|
|
| def get_context_templates(model, tok, length_params, device): |
| global CONTEXT_TEMPLATES_CACHE |
|
|
| if CONTEXT_TEMPLATES_CACHE is None: |
| CONTEXT_TEMPLATES_CACHE = [] |
| prompt_tok = tok( |
| ["I", "You", "Because", 'Yes', 'Q: '], |
| padding=True, |
| return_tensors="pt" |
| ).to(device) |
| for length, n_gen in length_params: |
|
|
| gen_token = model.generate( |
| input_ids=prompt_tok['input_ids'], |
| attention_mask=prompt_tok['attention_mask'], |
| max_new_tokens=length, |
| num_beams=n_gen // 5, |
| num_return_sequences=n_gen // 5, |
| pad_token_id=tok.eos_token_id, |
| ) |
| CONTEXT_TEMPLATES_CACHE += tok.batch_decode(gen_token, skip_special_tokens=True) |
| CONTEXT_TEMPLATES_CACHE = ['{}'] + [_ + ' {}' for _ in CONTEXT_TEMPLATES_CACHE] |
| return CONTEXT_TEMPLATES_CACHE |
|
|
|
|
