| import transformers |
| import torch |
| import os |
| import numpy as np |
| import datetime |
| import struct |
| from torch.nn.utils.rnn import pad_sequence |
| import torch.nn.functional as F |
|
|
| 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 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 |
|
|
| def tokenize(batch, tokenizer, device, test=False): |
| prompt, label = batch["prompt"], batch["target_new"] |
| if not isinstance(prompt, list): |
| prompt=[prompt] |
| if not isinstance(label, list): |
| label=[label] |
| mask_token = -100 |
| if test or not label: |
| tokens = tokenizer(list(prompt), return_tensors="pt", padding=True, truncation=True) |
| tokens["labels"] = tokens["input_ids"].clone() |
| tokens["labels"][tokens["input_ids"] == tokenizer.pad_token_id] = mask_token |
|
|
| else: |
| full_prompt = [f"{p} {l}" for p, l in zip(prompt, label)] |
| prompt_ids = tokenizer(list(prompt), return_tensors="pt", padding=True, truncation=True)["input_ids"] |
| num_prompt_toks = [int((i != tokenizer.pad_token_id).sum()) for i in prompt_ids] |
| tokens = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True) |
| tokens["labels"] = tokens["input_ids"].clone() |
| for i in range(len(prompt)): |
| tokens["labels"][i][:num_prompt_toks[i]] = mask_token |
|
|
| tokens["labels"][tokens["input_ids"] == tokenizer.pad_token_id] = mask_token |
| |
| tokens = {f"{k1}" : v1.to(device) for k1, v1 in tokens.items()} |
| return tokens |
|
|
|
|
