File size: 5,336 Bytes
b0e8c81 59d55b8 b0e8c81 59d55b8 b0e8c81 59d55b8 b0e8c81 |
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 |
import torch
from transformers import AutoTokenizer, PretrainedConfig, T5Config, PreTrainedModel, T5ForConditionalGeneration, \
AutoModelForSeq2SeqLM, Adafactor
from typing import Optional, List, Callable, Mapping, Any, Union
import os
class STEPFinetuningModelConfig(T5Config):
model_type = "STEP_finetuning"
def __init__(self,
num_examples: int = 512,
prefix_length: int = 10,
random_selection: bool = True,
# don't change these unless you change what the prefix of the model is initialized with:
prefix_max_init_length: int = 20,
num_precomputed_examples: int = 1000,
**kwargs):
# These are all about the initialization of the prefix.
self.num_examples = num_examples
self.prefix_length = prefix_length
self.random_selection = random_selection
self.prefix_max_init_length = prefix_max_init_length
self.num_precomputed_examples = num_precomputed_examples
super().__init__(**kwargs)
class STEPFinetuningModel(PreTrainedModel):
config_class = STEPFinetuningModelConfig
def __init__(self, config: STEPFinetuningModelConfig):
super().__init__(config)
self.model = T5ForConditionalGeneration(config)
# Initialize the prefix with NaNs.
self.register_buffer("prefix_init_tensor", torch.zeros(config.num_precomputed_examples, config.prefix_max_init_length, config.d_model))
# There are two cases: (1) we initialize the model after STEP-pretraining, i.e. the tunable prefix is not set
# and (2) the model has been fine-tuned on downstream data, and hence there is meaningful data in the tunable prefix
# Initialize the prefix with NaNs. If we initialize from STEP-pretraining, this will be overwritten by a custom version of from_pretrained
# if we initialize after fine-tuning, the NaNs will be overwritten anyway.
self.prefix_embedding = torch.nn.Parameter(torch.nan + torch.zeros((1, self.config.prefix_length, self.config.d_model)))
self.prefix_has_been_initialized = False
def _initialize_prefix(self):
prefix_init_tensor = self.prefix_init_tensor
if self.config.random_selection:
# randomize selection of edgewise tranformations to average for initialization the prefix.
prefix_init_tensor = prefix_init_tensor[torch.randperm(prefix_init_tensor.shape[0]), :, :]
prefix_init_tensor = prefix_init_tensor[:self.config.num_examples, :self.config.prefix_length,
:] # shape (num ex, prefix length, d model)
self.prefix_embedding.data.copy_(prefix_init_tensor.mean(dim=0, keepdims=True))
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
*model_args,
**kwargs,
):
model = super(STEPFinetuningModel, cls).from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
if torch.all(model.prefix_embedding.isnan()):
model._initialize_prefix()
return model
def prepare_input(self, kwargs):
"""
Prepends the prefix to the given input.
:param kwargs:
:return:
"""
input_ids = kwargs["input_ids"]
embedded_inputs = self.model.get_input_embeddings()(input_ids)
batch_size = input_ids.shape[0]
prefix = torch.repeat_interleave(self.prefix_embedding, batch_size, 0) #shape (batch, prefix length, embed dim)
kwargs = dict(kwargs)
embedded_inputs = torch.cat([prefix, embedded_inputs], dim=1) # shape (batch, prefix + seq length, embed dim)
del kwargs["input_ids"]
kwargs["inputs_embeds"] = embedded_inputs
if "attention_mask" in kwargs:
ones = torch.ones((batch_size, self.config.prefix_length), device=embedded_inputs.device, dtype=kwargs["attention_mask"].dtype)
input_mask = torch.cat([ones, kwargs["attention_mask"]], dim=1)
kwargs["attention_mask"] = input_mask
return kwargs
def forward(self, **kwargs):
return self.model(**self.prepare_input(kwargs))
def generate(self, **kwargs):
return self.model.generate(**self.prepare_input(kwargs))
def get_optimizer(self, optimizer: Callable[..., torch.optim.Optimizer] = None, prefix_lr:float = 10.0, **kwargs) -> torch.optim.Optimizer:
"""
Return an optimizer that uses a different learning rate (typically higher) for the prefix than for the rest of the model.
"""
prefix_params = []
other_params = []
for name, param in self.named_parameters():
if name == "prefix_embedding":
prefix_params.append(param)
else:
other_params.append(param)
if optimizer is None:
# The optimizer used in the paper.
hparams = {"scale_parameter": False, "relative_step": False, "warmup_init": False, "lr": 1e-4}
return Adafactor(params=[{"params": prefix_params, "lr": prefix_lr}, {"params": other_params}], **(hparams | kwargs))
return optimizer(params=[{"params": prefix_params, "lr": prefix_lr}, {"params": other_params}], **kwargs)
|