diff --git a/README.md b/README.md index fe2639a..dde65c9 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # 🤗 pets Parameter-Efficient Tuning at Scale with 🤗 Accelerate -Supported moethods: +Supported methods: 1. Prefix Tuning 2. P-Tuning 3. Prompt Tuning -4. LoRA [in progress] +4. LoRA [in backlog] ## Models support matrix @@ -25,18 +25,17 @@ Supported moethods: ### Causal Language Modeling | | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | | --------- | ---- | ---- | ---- | ---- | -| GPT-2 | | | | | -| Bloom | | | | | -| OPT | | | | | -| GPT-Neo | | | | | -| GPT-J | | | | | -| BART | | | | | +| GPT-2 | ✅ | ✅ | ✅ | | +| Bloom | ✅ | ✅ | ✅ | | +| OPT | ✅ | ✅ | ✅ | | +| GPT-Neo | ✅ | ✅ | ✅ | | +| GPT-J | ✅ | ✅ | ✅ | | ### Conditional Generation | | Prefix Tuning | P-Tuning | Prompt Tuning | LoRA | | --------- | ---- | ---- | ---- | ---- | -| T5 | | | | | -| BART | | | | | +| T5 | ✅ | ✅ | ✅ | | +| BART | ✅ | ✅ | ✅ | | diff --git a/src/pet/__init__.py b/src/pet/__init__.py index 9598427..77a32cc 100644 --- a/src/pet/__init__.py +++ b/src/pet/__init__.py @@ -9,12 +9,18 @@ from .pet_model import ( PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification, - PromptEncoderType, + PETPluginBase, + PETType, ) +from .task_mapping import MODEL_TYPE_TO_PROMPT_MODEL_MAPPING from .tuners import ( PrefixEncoder, + PrefixTuningConfig, PromptEmbedding, PromptEncoder, + PromptEncoderConfig, PromptEncoderReparameterizationType, + PromptTuningConfig, PromptTuningInit, ) +from .utils import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py index 814473c..45a6264 100644 --- a/src/pet/pet_model.py +++ b/src/pet/pet_model.py @@ -1,4 +1,3 @@ -import enum import inspect import warnings @@ -8,98 +7,78 @@ from transformers import PreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput from .tuners import PrefixEncoder, PromptEmbedding, PromptEncoder - - -class PromptEncoderType(str, enum.Enum): - PROMPT_TUNING = "PROMPT_TUNING" - P_TUNING = "P_TUNING" - PREFIX_TUNING = "PREFIX_TUNING" - LORA = "LORA" +from .utils import PETConfig, PETType, TaskType class PETModel(torch.nn.Module): - def __init__(self, model): + def __init__(self, model, pet_config: PETConfig): super().__init__() self.model = model - self.prompt_learning_config = model.config.prompt_learning_config + self.pet_config = pet_config num_transformer_submodules = 0 transformer_backbone = None for name, module in self.model.named_children(): if isinstance(module, PreTrainedModel): + # Make sure to freeze Tranformers model + for param in module.parameters(): + param.requires_grad = False if transformer_backbone is None: transformer_backbone = module self.transformer_backbone_name = name num_transformer_submodules += 1 - self.prompt_learning_config["num_transformer_submodules"] = num_transformer_submodules + self.pet_config.num_transformer_submodules = 2 if self.pet_config.task_type == TaskType.SEQ_2_SEQ_LM else 1 for named_param, value in list(transformer_backbone.named_parameters()): if value.shape[0] == model.config.vocab_size: self.word_embeddings = transformer_backbone.get_submodule(named_param.replace(".weight", "")) break - # Make sure to freeze Tranformers model - for param in transformer_backbone.parameters(): - param.requires_grad = False - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING: - prompt_encoder = PromptEncoder(self.prompt_learning_config) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: - prompt_encoder = PrefixEncoder(self.prompt_learning_config) + if self.pet_config.pet_type == PETType.PROMPT_TUNING: + prompt_encoder = PromptEmbedding(self.pet_config, self.word_embeddings) + elif self.pet_config.pet_type == PETType.P_TUNING: + prompt_encoder = PromptEncoder(self.pet_config) + elif self.pet_config.pet_type == PETType.PREFIX_TUNING: + prompt_encoder = PrefixEncoder(self.pet_config) else: raise ValueError("Not supported") self.prompt_encoder = prompt_encoder self.prompt_tokens = torch.arange( - self.prompt_learning_config["num_virtual_tokens"] - * self.prompt_learning_config["num_transformer_submodules"] + self.pet_config.num_virtual_tokens * self.pet_config.num_transformer_submodules ).long() def get_prompt(self, batch_size): prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.model.device) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: - prompt_tokens = prompt_tokens[:, : self.prompt_learning_config["num_virtual_tokens"]] - if self.prompt_learning_config.get("inference_mode", False): + if self.pet_config.pet_type == PETType.PREFIX_TUNING: + prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens] + if self.pet_config.inference_mode: past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: past_key_values = self.prompt_encoder(prompt_tokens) past_key_values = past_key_values.view( batch_size, - self.prompt_learning_config["num_virtual_tokens"], - self.prompt_learning_config["num_layers"] - * self.prompt_learning_config["num_transformer_submodules"] - * 2, - self.prompt_learning_config["num_attention_heads"], - self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], + self.pet_config.num_virtual_tokens, + self.pet_config.num_layers * 2, + self.pet_config.num_attention_heads, + self.pet_config.token_dim // self.pet_config.num_attention_heads, ) + if self.pet_config.num_transformer_submodules == 2: + past_key_values = torch.cat([past_key_values, past_key_values], dim=2) past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( - self.prompt_learning_config["num_transformer_submodules"] * 2 + self.pet_config.num_transformer_submodules * 2 ) - if "postprocess_past_key_value_function" in self.prompt_learning_config["prompt_encoder_config"]: - post_process_fn = self.prompt_learning_config["prompt_encoder_config"][ - "postprocess_past_key_value_function" - ] + if self.pet_config.postprocess_past_key_value_function is not None: + post_process_fn = self.pet_config.postprocess_past_key_value_function past_key_values = post_process_fn(past_key_values) return past_key_values else: - if self.prompt_learning_config.get("inference_mode", False): + if self.pet_config.inference_mode: prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) else: prompts = self.prompt_encoder(prompt_tokens) return prompts - -class PETModelForSequenceClassification(PETModel): - def __init__(self, model): - super().__init__(model) - self.config = self.model.config - - for name, module in self.model.named_children(): - if isinstance(module, torch.nn.Linear): - self.cls_layer_name = name - break - + def print_trainable_parameters(self): trainable_params = 0 all_param = 0 for _, param in self.named_parameters(): @@ -110,6 +89,17 @@ class PETModelForSequenceClassification(PETModel): f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" ) + +class PETModelForSequenceClassification(PETModel): + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) + self.config = self.model.config + + for name, module in self.model.named_children(): + if isinstance(module, torch.nn.Linear): + self.cls_layer_name = name + break + def forward( self, input_ids=None, @@ -124,11 +114,9 @@ class PETModelForSequenceClassification(PETModel): return_dict = return_dict if return_dict is not None else self.config.use_return_dict batch_size = input_ids.shape[0] - if attention_mask is not None and self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + if attention_mask is not None and self.pet_config.pet_type != PETType.LORA: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.model.device - ) + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.model.device) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -143,15 +131,13 @@ class PETModelForSequenceClassification(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.pet_config.pet_type == PETType.PREFIX_TUNING: return self.prefix_tuning_forward(input_ids=input_ids, **kwargs) else: if kwargs.get("token_type_ids", None) is not None: kwargs["token_type_ids"] = torch.cat( ( - torch.zeros(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.model.device - ), + torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.model.device), kwargs["token_type_ids"], ), dim=1, @@ -235,20 +221,10 @@ class PETModelForSequenceClassification(PETModel): class PETModelForCausalLM(PETModel): - def __init__(self, model): - super().__init__(model) + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) self.config = self.model.config - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - def forward( self, input_ids=None, @@ -261,10 +237,10 @@ class PETModelForCausalLM(PETModel): **kwargs, ): batch_size = input_ids.shape[0] - if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: + if self.pet_config.pet_type != PETType.LORA: if attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.model.device ) attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) @@ -285,38 +261,29 @@ class PETModelForCausalLM(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.pet_config.pet_type == PETType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.model(input_ids=input_ids, past_key_values=past_key_values, **kwargs) else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) - # concat prompt labels - if kwargs["labels"] is not None: - prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.model.device - ) - kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) + if self.pet_config.pet_type != PETType.LORA: + # concat prompt labels + if labels is not None: + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to( + self.model.device + ) + kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1) return self.model(inputs_embeds=inputs_embeds, **kwargs) class PETModelForSeq2SeqLM(PETModel): - def __init__(self, model): - super().__init__(model) + def __init__(self, model, pet_config: PETConfig): + super().__init__(model, pet_config) self.config = self.model.config - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - def forward( self, input_ids=None, @@ -333,22 +300,13 @@ class PETModelForSeq2SeqLM(PETModel): ): batch_size = input_ids.shape[0] - if self.prompt_learning_config["prompt_encoder_type"] != PromptEncoderType.LORA: - if attention_mask is not None: + if self.pet_config.pet_type != PETType.LORA: + if decoder_attention_mask is not None: # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( self.model.device ) - attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) - if decoder_attention_mask is not None: - decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) - - # concat prompt labels - if labels is not None: - prefix_labels = torch.full((batch_size, self.prompt_learning_config["num_virtual_tokens"]), -100).to( - self.model.device - ) - labels = torch.cat((prefix_labels, labels), dim=1) + decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1) if kwargs.get("position_ids", None) is not None: warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.") @@ -367,7 +325,7 @@ class PETModelForSeq2SeqLM(PETModel): } ) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PREFIX_TUNING: + if self.pet_config.pet_type == PETType.PREFIX_TUNING: past_key_values = self.get_prompt(batch_size) return self.model( input_ids=input_ids, decoder_input_ids=decoder_input_ids, past_key_values=past_key_values, **kwargs @@ -375,13 +333,30 @@ class PETModelForSeq2SeqLM(PETModel): else: if inputs_embeds is None: inputs_embeds = self.word_embeddings(input_ids) - if decoder_inputs_embeds is None: + if decoder_inputs_embeds is None and decoder_input_ids is None: + from transformers.models.bart.modeling_bart import shift_tokens_right + + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) decoder_inputs_embeds = self.word_embeddings(decoder_input_ids) + + if self.pet_config.pet_type != PETType.LORA: + if attention_mask is not None: + # concat prompt attention mask + prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to( + self.model.device + ) + kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1) + # concat prompt labels + if labels is not None: + prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to( + self.model.device + ) + kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1) prompts = self.get_prompt(batch_size=batch_size) - inputs_embeds = torch.cat( - (prompts[:, : self.prompt_learning_config["num_virtual_tokens"]], inputs_embeds), dim=1 - ) + inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1) decoder_inputs_embeds = torch.cat( - (prompts[:, self.prompt_learning_config["num_virtual_tokens"] :], decoder_inputs_embeds), dim=1 + (prompts[:, self.pet_config.num_virtual_tokens :], decoder_inputs_embeds), dim=1 ) return self.model(inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, **kwargs) diff --git a/src/pet/prompt_learning_legacy.py b/src/pet/prompt_learning_legacy.py deleted file mode 100644 index b08d018..0000000 --- a/src/pet/prompt_learning_legacy.py +++ /dev/null @@ -1,657 +0,0 @@ -import enum -import functools -import math -import os -from collections import OrderedDict - -import torch -from accelerate import Accelerator -from accelerate.state import AcceleratorState -from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin -from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from torch.utils.data import DataLoader -from transformers import ( - AutoModelForSequenceClassification, - AutoTokenizer, - PreTrainedModel, - get_linear_schedule_with_warmup, - set_seed, -) -from transformers.modeling_outputs import SequenceClassifierOutput - -import evaluate -from datasets import load_dataset - - -class PromptEncoderReparameterizationType(str, enum.Enum): - MLP = "MLP" - LSTM = "LSTM" - - -class PromptEncoderType(str, enum.Enum): - PROMPT_TUNING = "PROMPT_TUNING" - P_TUNING_V1 = "P_TUNING_V1" - P_TUNING_V2 = "P_TUNING_V2" - - -class PromptTuningInit(str, enum.Enum): - TEXT = "TEXT" - RANDOM = "RANDOM" - - -class PromptEncoder(torch.nn.Module): - """ - The prompt encoder network that is used to generate the virtual token embeddings for p-tuning. - """ - - def __init__(self, config): - super().__init__() - self.token_dim = config["token_dim"] - self.input_size = config["token_dim"] - self.output_size = config["token_dim"] - self.hidden_size = config["prompt_hidden_size"] - self.total_virtual_tokens = config["num_virtual_tokens"] - self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] - - # embedding - self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) - if not config.get("inference_mode", False): - if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - if "dropout" not in config["prompt_encoder_config"]: - lstm_dropout = 0.0 - else: - lstm_dropout = config["prompt_encoder_config"]["dropout"] - - if "num_layers" not in config["prompt_encoder_config"]: - num_layers = 2 - else: - num_layers = config["prompt_encoder_config"]["num_layers"] - # LSTM - self.lstm_head = torch.nn.LSTM( - input_size=self.input_size, - hidden_size=self.hidden_size, - num_layers=num_layers, - dropout=lstm_dropout, - bidirectional=True, - batch_first=True, - ) - - self.mlp_head = torch.nn.Sequential( - torch.nn.Linear(self.hidden_size * 2, self.hidden_size * 2), - torch.nn.ReLU(), - torch.nn.Linear(self.hidden_size * 2, self.output_size), - ) - - elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - layers = [ - torch.nn.Linear(self.input_size, self.hidden_size), - torch.nn.ReLU(), - ] - layers.extend( - [ - torch.nn.Linear(self.hidden_size, self.hidden_size), - torch.nn.ReLU(), - ] - ) - layers.append(torch.nn.Linear(self.hidden_size, self.output_size)) - self.mlp_head = torch.nn.Sequential(*layers) - - else: - raise ValueError( - "Prompt encoder type not recognized. " " Please use one of MLP (recommended) or LSTM." - ) - - def forward(self, indices): - input_embeds = self.embedding(indices) - if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - output_embeds = self.mlp_head(self.lstm_head(input_embeds)[0]) - elif self.encoder_type == PromptEncoderReparameterizationType.MLP: - output_embeds = self.mlp_head(input_embeds) - else: - raise ValueError("Prompt encoder type not recognized. Please use one of MLP (recommended) or LSTM.") - - return output_embeds - - -class PrefixEncoder(torch.nn.Module): - r""" - The torch.nn model to encode the prefix - - Input shape: (batch-size, prefix-length) - - Output shape: (batch-size, prefix-length, 2*layers*hidden) - """ - - def __init__(self, config): - super().__init__() - self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] - if self.prefix_projection and not config.get("inference_mode", False): - # Use a two-layer MLP to encode the prefix - self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) - self.trans = torch.nn.Sequential( - torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), - torch.nn.Tanh(), - torch.nn.Linear( - config["prompt_hidden_size"], - config["num_layers"] * 2 * config["token_dim"], - ), - ) - else: - self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], - config["num_layers"] * 2 * config["token_dim"], - ) - - def forward(self, prefix: torch.Tensor): - if self.prefix_projection: - prefix_tokens = self.embedding(prefix) - past_key_values = self.trans(prefix_tokens) - else: - past_key_values = self.embedding(prefix) - return past_key_values - - -class PromptEmbedding(torch.nn.Module): - def __init__(self, config, word_embeddings): - super().__init__() - - total_virtual_tokens = config["num_virtual_tokens"] - self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) - if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: - from transformers import AutoTokenizer - - self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) - self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] - init_token_ids = self.tokenizer(self.init_text)["input_ids"] - # Trim or iterate until num_text_tokens matches total_virtual_tokens - num_text_tokens = len(init_token_ids) - if num_text_tokens > total_virtual_tokens: - init_token_ids = init_token_ids[:total_virtual_tokens] - elif num_text_tokens < total_virtual_tokens: - num_reps = math.ceil(total_virtual_tokens / num_text_tokens) - init_token_ids = init_token_ids * num_reps - init_token_ids = init_token_ids[:total_virtual_tokens] - - word_embedding_weights = word_embeddings(torch.LongTensor(init_token_ids)).detach().clone() - self.embedding.weight = torch.nn.Parameter(word_embedding_weights) - - def forward(self, indices): - # Just get embeddings and dropout - prompt_embeddings = self.embedding(indices) - return prompt_embeddings - - -class PromptModel(torch.nn.Module): - def __init__(self, model): - super().__init__() - self.prompt_learning_config = model.config.prompt_learning_config - - modules = list(model._modules) - - for module in modules: - if isinstance(model.get_submodule(module), PreTrainedModel): - self.transformer_backbone = model.get_submodule(module) - break - - for named_param, value in list(self.transformer_backbone.named_parameters()): - if value.shape[0] == model.config.vocab_size: - self.word_embeddings = self.transformer_backbone.get_submodule(named_param.replace(".weight", "")) - break - - # Make sure to freeze Tranformers model - for param in self.transformer_backbone.parameters(): - param.requires_grad = False - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.PROMPT_TUNING: - prompt_encoder = PromptEmbedding(self.prompt_learning_config, self.word_embeddings) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V1: - prompt_encoder = PromptEncoder(self.prompt_learning_config) - elif self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - prompt_encoder = PrefixEncoder(self.prompt_learning_config) - else: - raise ValueError("Not supported") - self.prompt_encoder = prompt_encoder - self.prompt_tokens = torch.arange(self.prompt_learning_config["num_virtual_tokens"]).long() - - def get_prompt(self, batch_size): - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.transformer_backbone.device) - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - if self.prompt_learning_config.get("inference_mode", False): - past_key_values = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) - else: - past_key_values = self.prompt_encoder(prompt_tokens) - past_key_values = past_key_values.view( - batch_size, - self.prompt_learning_config["num_virtual_tokens"], - self.prompt_learning_config["num_layers"] * 2, - self.prompt_learning_config["num_attention_heads"], - self.prompt_learning_config["token_dim"] // self.prompt_learning_config["num_attention_heads"], - ) - past_key_values = self.dropout(past_key_values) - past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2) - return past_key_values - else: - if self.prompt_learning_config.get("inference_mode", False): - prompts = self.prompt_encoder.embedding.weight.repeat(batch_size, 1, 1) - else: - prompts = self.prompt_encoder(prompt_tokens) - return prompts - - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.transformer_backbone.device) - prompt_embeddings = self.prompt_encoder(prompt_tokens).detach().cpu() - if destination is None: - state_dict_ = OrderedDict() - else: - state_dict_ = destination - state_dict_["prompt_embeddings"] = prompt_embeddings[0] - return state_dict_ - - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - self.prompt_encoder.embedding.load_state_dict({"weight": state_dict["prompt_embeddings"]}, strict) - - -class PromptModelForSequenceClassification(PromptModel): - def __init__(self, model): - super().__init__(model) - if "dropout" in [name for name, _ in model.named_children()]: - self.dropout = model.dropout - else: - self.dropout = torch.nn.Dropout(model.config.hidden_dropout_prob) - self.classifier = model.classifier - self.num_labels = model.num_labels - self.config = model.config - self.modules_to_save = ("prompt_encoder", "classifier") - - trainable_params = 0 - all_param = 0 - for _, param in self.named_parameters(): - all_param += param.numel() - if param.requires_grad: - trainable_params += param.numel() - print( - f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}" - ) - - def forward( - self, - input_ids=None, - attention_mask=None, - inputs_embeds=None, - labels=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - **kwargs, - ): - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - batch_size = input_ids.shape[0] - # concat prompt attention mask - prefix_attention_mask = torch.ones(batch_size, self.prompt_learning_config["num_virtual_tokens"]).to( - self.transformer_backbone.device - ) - attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1) - - if self.prompt_learning_config["prompt_encoder_type"] == PromptEncoderType.P_TUNING_V2: - past_key_values = self.get_prompt(batch_size=batch_size) - - outputs = self.transformer_backbone( - input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - past_key_values=past_key_values, - **kwargs, - ) - - pooled_output = outputs[1] if len(outputs) > 1 else outputs[0] - else: - raw_embedding = self.word_embeddings(input_ids) - prompts = self.get_prompt(batch_size=batch_size) - inputs_embeds = torch.cat((prompts, raw_embedding), dim=1) - - outputs = self.transformer_backbone( - # input_ids, - attention_mask=attention_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - # **kwargs, - # past_key_values=past_key_values, - ) - - sequence_output = outputs[0] - sequence_output = sequence_output[:, self.prompt_learning_config["num_virtual_tokens"] :, :].contiguous() - pooled_output = sequence_output[:, 0] - - if ( - "pooler" in [name for name, _ in self.transformer_backbone.named_children()] - and self.transformer_backbone.pooler is not None - ): - pooled_output = self.transformer_backbone.pooler.dense(pooled_output) - pooled_output = self.transformer_backbone.pooler.activation(pooled_output) - - pooled_output = self.dropout(pooled_output) - logits = self.classifier(pooled_output) - - loss = None - if labels is not None: - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - def state_dict(self, destination=None, prefix=None, keep_vars=False): - """ - No frozen model parameters are stored in the state dict. - """ - if destination is None: - state_dict_ = OrderedDict() - else: - state_dict_ = destination - state_dict_["prompt_encoder"] = super().state_dict() - state_dict_["classifier"] = self.classifier.state_dict() - if AcceleratorState().fsdp_plugin is not None: - state_dict_["_flat_param"] = None - return state_dict_ - - def load_state_dict(self, state_dict, strict: bool = True): - """ - Custom load state dict method that only loads prompt table and prompt encoder parameters. Matching load method - for this class' custom state dict method. - """ - super().load_state_dict(state_dict["prompt_encoder"], strict) - self.classifier.load_state_dict(state_dict["classifier"], strict) - - def clean_state_dict(self, state_dict): - if AcceleratorState().fsdp_plugin is not None: - new_state_dict = OrderedDict() - for key in self.modules_to_save: - new_state_dict[key] = state_dict[key].copy() - state_dict = new_state_dict - return state_dict - - -model_type_to_prompt_model_mapping = {"SequenceClassification": PromptModelForSequenceClassification} -num_virtual_tokens = 30 -model_name_or_path = "roberta-large" -tokenizer_name_or_path = "roberta-large" - -prompt_tuning_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "PROMPT_TUNING", - "prompt_encoder_config": { - "prompt_tuning_init": "TEXT", - "tokenizer_name_or_path": tokenizer_name_or_path, - "prompt_tuning_text": "Output is true or false. Task requires to recognize" - " whether the meaning of one text is entailed (can be inferred) from the other text.", - }, -} - - -p_tuning_v1_mlp_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V1", - "prompt_encoder_config": {"prompt_reparam_type": "MLP"}, -} - -p_tuning_v1_lstm_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V1", - "prompt_encoder_config": {"prompt_reparam_type": "LSTM"}, -} - -p_tuning_v2_no_proj_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V2", - "prompt_encoder_config": {"prefix_projection": False}, -} - -p_tuning_v2_proj_config = { - "num_virtual_tokens": num_virtual_tokens, - "prompt_encoder_type": "P_TUNING_V2", - "prompt_encoder_config": {"prefix_projection": True}, -} - - -def prepare_prompt_model(model, prompt_learning_config): - config = model.config.to_dict() - if "num_layers" not in prompt_learning_config: - if "num_hidden_layers" in config: - num_layers = config["num_hidden_layers"] - elif "num_layers" in config: - num_layers = config["num_layers"] - else: - raise ValueError("Please specify `num_layers` in `prompt_learning_config`") - prompt_learning_config["num_layers"] = num_layers - - if "token_dim" not in prompt_learning_config: - if "hidden_size" in config: - token_dim = config["hidden_size"] - elif "n_embd" in config: - token_dim = config["n_embd"] - elif "d_model" in config: - token_dim = config["d_model"] - else: - raise ValueError("Please specify `token_dim` in `prompt_learning_config`") - prompt_learning_config["token_dim"] = token_dim - - if "num_attention_heads" not in prompt_learning_config: - if "num_attention_heads" in config: - num_attention_heads = config["num_attention_heads"] - elif "n_head" in config: - num_attention_heads = config["n_head"] - elif "num_heads" in config: - num_attention_heads = config["num_heads"] - else: - raise ValueError("Please specify `num_attention_heads` in `prompt_learning_config`") - prompt_learning_config["num_attention_heads"] = num_attention_heads - - if "prompt_hidden_size" not in prompt_learning_config: - prompt_learning_config["prompt_hidden_size"] = token_dim - - model.config.prompt_learning_config = prompt_learning_config - model_type = model.__class__.__name__.split("For") - if len(model_type) < 2: - raise ValueError("Model Type not supported") - model_cls = model_type_to_prompt_model_mapping[model_type[1]] - prompt_model = model_cls(model) - return prompt_model - - -def fsdp_auto_wrap_policy(model): - def wrap_layers_with_required_grads(module): - if ( - len(list(module.children())) == 0 - and len(list(module.named_parameters())) > 0 - and module.weight.requires_grad - ): - return True - return False - - transformer_cls_to_wrap = { - PrefixEncoder, - PromptEmbedding, - PromptEncoder, - PromptModel, - FullyShardedDataParallelPlugin.get_module_class_from_name( - model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "") - ), - } - policy_1 = functools.partial( - transformer_auto_wrap_policy, - transformer_layer_cls=transformer_cls_to_wrap, - ) - policy_2 = functools.partial( - lambda_auto_wrap_policy, - lambda_fn=wrap_layers_with_required_grads, - ) - auto_wrap_policy = functools.partial(_or_policy, policies=[policy_1, policy_2]) - return auto_wrap_policy - - -def main(): - accelerator = Accelerator() - task = "rte" - batch_size = 16 - lr = 5e-3 - num_epochs = 100 - # device = "cuda" - seed = 11 - set_seed(seed) - - model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path) - model = prepare_prompt_model( - model, p_tuning_v2_no_proj_config - ) # p_tuning_v2_proj_config)#p_tuning_v2_no_proj_config) - # model = model.to("cuda") - - tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) - datasets = load_dataset("glue", task) - metric = evaluate.load("glue", task) - - def tokenize_function(examples): - # max_length=None => use the model max length (it's actually the default) - outputs = tokenizer( - examples["sentence1"], - examples["sentence2"], - truncation=True, - max_length=None, - ) - return outputs - - # Apply the method we just defined to all the examples in all the splits of the dataset - # starting with the main process first: - tokenized_datasets = datasets.map( - tokenize_function, - batched=True, - remove_columns=["idx", "sentence1", "sentence2"], - ) - - # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the - # transformers library - tokenized_datasets = tokenized_datasets.rename_column("label", "labels") - - def collate_fn(examples): - return tokenizer.pad(examples, padding="longest", return_tensors="pt") - - # Instantiate dataloaders. - train_dataloader = DataLoader( - tokenized_datasets["train"], - shuffle=True, - collate_fn=collate_fn, - batch_size=batch_size, - ) - eval_dataloader = DataLoader( - tokenized_datasets["validation"], - shuffle=False, - collate_fn=collate_fn, - batch_size=batch_size, - ) - - # Instantiate optimizer - optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr) - - # Instantiate scheduler - lr_scheduler = get_linear_schedule_with_warmup( - optimizer=optimizer, - num_warmup_steps=0, - num_training_steps=(len(train_dataloader) * num_epochs), - ) - - accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model) - - ( - model, - train_dataloader, - eval_dataloader, - optimizer, - lr_scheduler, - ) = accelerator.prepare(model, train_dataloader, eval_dataloader, optimizer, lr_scheduler) - accelerator.print(model) - - for epoch in range(num_epochs): - model.train() - total_loss = 0 - for step, batch in enumerate(train_dataloader): - # batch.to(device) - outputs = model(**batch) - loss = outputs.loss - total_loss += loss.detach().float() - loss.backward() - optimizer.step() - lr_scheduler.step() - optimizer.zero_grad() - - model.eval() - for step, batch in enumerate(eval_dataloader): - # batch.to(device) - with torch.no_grad(): - outputs = model(**batch) - predictions = outputs.logits.argmax(dim=-1) - predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"])) - metric.add_batch( - predictions=predictions, - references=references, - ) - - eval_metric = metric.compute() - accelerator.print(f"epoch {epoch}:", eval_metric) - accelerator.print(f"epoch {epoch} train loss:", total_loss / len(train_dataloader)) - - from torch.distributed.fsdp.fully_sharded_data_parallel import FullStateDictConfig - from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP - from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType - - FSDP.set_state_dict_type( - model, - StateDictType.FULL_STATE_DICT, - FullStateDictConfig(offload_to_cpu=True, rank0_only=True), - ) - state_dict = model.state_dict() - state_dict = model.clean_state_dict(state_dict) - accelerator.print(state_dict) - - torch.save(state_dict, "p_tuning_v2.pt") - - -if __name__ == "__main__": - main() diff --git a/src/pet/task_mapping.py b/src/pet/task_mapping.py new file mode 100644 index 0000000..00e7415 --- /dev/null +++ b/src/pet/task_mapping.py @@ -0,0 +1,17 @@ +from .pet_model import PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification +from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig +from .utils import PETConfig + + +MODEL_TYPE_TO_PROMPT_MODEL_MAPPING = { + "SEQ_CLS": PETModelForSequenceClassification, + "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM, + "CAUSAL_LM": PETModelForCausalLM, +} + +PET_TYPE_TO_CONFIG_MAPPING = { + "PROMPT_TUNING": PromptTuningConfig, + "PREFIX_TUNING": PrefixTuningConfig, + "P_TUNING": PromptEncoderConfig, + "LORA": PETConfig, +} diff --git a/src/pet/tuners/__init__.py b/src/pet/tuners/__init__.py index 41ae36c..79fcc15 100644 --- a/src/pet/tuners/__init__.py +++ b/src/pet/tuners/__init__.py @@ -2,6 +2,6 @@ # There's no way to ignore "F401 '...' imported but unused" warnings in this # module, but to preserve other warnings. So, don't check this module at all -from .p_tuning import PromptEncoder, PromptEncoderReparameterizationType -from .prefix_tuning import PrefixEncoder -from .prompt_tuning import PromptEmbedding, PromptTuningInit +from .p_tuning import PromptEncoder, PromptEncoderConfig, PromptEncoderReparameterizationType +from .prefix_tuning import PrefixEncoder, PrefixTuningConfig +from .prompt_tuning import PromptEmbedding, PromptTuningConfig, PromptTuningInit diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py index 044a482..7fbb0b4 100644 --- a/src/pet/tuners/lora.py +++ b/src/pet/tuners/lora.py @@ -1 +1,39 @@ # todo +import torch +from transformers import Conv1D + +import loralib as lora + + +class LoRAModel(torch.nn.Module): + def __init__(self, config, model): + super().__init__() + self.config = config + self.model = model + + def find_and_replace(self): + key_list = [key for key, _ in self.model.named_modules()] + for key in key_list: + if any(key.endswith(target_key) for target_key in self.config["target_module_keys"]): + parent, target_name, target = self.get_submodules(key) + if isinstance(target, torch.nn.Linear): + new_module = lora.Linear( + target.in_features, target.out_features, **self.config["prompt_encoder_config"] + ) + elif isinstance(target, torch.nn.Conv1d, Conv1D): + new_module = lora.LoRAConv1d( + target.in_channels, target.out_channels, target.kernel_size, bias=target.bias is not None + ) + self.replace_module(parent, target_name, new_module) + + def get_submodules(self, key): + parent = self.model.get_submodule(".".join(key.split(".")[:-1])) + target_name = key.split(".")[:-1] + target = self.model.get_submodule(key) + return parent, target_name, target + + def replace_module(self, parent_module, child_name, new_module, old_module): + setattr(parent_module, child_name, new_module) + new_module.weight = old_module.weight.clone() + if old_module.bias is not None: + new_module.bias = old_module.bias.clone() diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 25149b9..8f8e030 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -1,13 +1,37 @@ import enum +from dataclasses import dataclass, field +from typing import Union import torch +from ..utils import PromptLearningConfig + class PromptEncoderReparameterizationType(str, enum.Enum): MLP = "MLP" LSTM = "LSTM" +@dataclass +class PromptEncoderConfig(PromptLearningConfig): + encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field( + default=PromptEncoderReparameterizationType.MLP, + metadata={"help": "How to reparameterize the prompt encoder"}, + ) + encoder_hidden_size: int = field( + default=256, + metadata={"help": "The hidden size of the prompt encoder reparameterization"}, + ) + encoder_num_layers: int = field( + default=2, + metadata={"help": "The number of layers of the prompt encoder reparameterization"}, + ) + encoder_dropout: float = field( + default=0.0, + metadata={"help": "The dropout of the prompt encoder reparameterization"}, + ) + + # Based on https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/nlp/modules/common/prompt_encoder.py # with some refactor class PromptEncoder(torch.nn.Module): @@ -17,26 +41,19 @@ class PromptEncoder(torch.nn.Module): def __init__(self, config): super().__init__() - self.token_dim = config["token_dim"] - self.input_size = config["token_dim"] - self.output_size = config["token_dim"] - self.hidden_size = config["prompt_hidden_size"] - self.total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] - self.encoder_type = config["prompt_encoder_config"]["prompt_reparam_type"] + self.token_dim = config.token_dim + self.input_size = self.token_dim + self.output_size = self.token_dim + self.hidden_size = config.encoder_hidden_size + self.total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules + self.encoder_type = config.encoder_reparameterization_type # embedding self.embedding = torch.nn.Embedding(self.total_virtual_tokens, self.token_dim) - if not config.get("inference_mode", False): + if not config.inference_mode: if self.encoder_type == PromptEncoderReparameterizationType.LSTM: - if "dropout" not in config["prompt_encoder_config"]: - lstm_dropout = 0.0 - else: - lstm_dropout = config["prompt_encoder_config"]["dropout"] - - if "num_layers" not in config["prompt_encoder_config"]: - num_layers = 2 - else: - num_layers = config["prompt_encoder_config"]["num_layers"] + lstm_dropout = config.encoder_dropout + num_layers = config.encoder_num_layers # LSTM self.lstm_head = torch.nn.LSTM( input_size=self.input_size, diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index 1319fe3..049138b 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -1,5 +1,26 @@ +from dataclasses import dataclass, field +from typing import Callable, Optional + import torch +from ..utils import PromptLearningConfig + + +@dataclass +class PrefixTuningConfig(PromptLearningConfig): + encoder_hidden_size: int = field( + default=256, + metadata={"help": "The hidden size of the encoder"}, + ) + prefix_projection: bool = field( + default=False, + metadata={"help": "Whether to project the prefix tokens"}, + ) + postprocess_past_key_value_function: Optional[Callable] = field( + default=None, + metadata={"help": "The function to postprocess the past key value"}, + ) + # Based on https://github.com/THUDM/P-tuning-v2/blob/main/model/prefix_encoder.py # with some refactor @@ -9,28 +30,26 @@ class PrefixEncoder(torch.nn.Module): Input shape: (batch_size, num_virtual_tokens) - Output shape: (batch_size, num_virtual_tokens, 2*(num_transformer_submodules)*layers*hidden) + Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden) """ def __init__(self, config): super().__init__() - self.prefix_projection = config["prompt_encoder_config"]["prefix_projection"] - if self.prefix_projection and not config.get("inference_mode", False): + self.prefix_projection = config.prefix_projection + token_dim = config.token_dim + num_layers = config.num_layers + encoder_hidden_size = config.encoder_hidden_size + num_virtual_tokens = config.num_virtual_tokens + if self.prefix_projection and not config.inference_mode: # Use a two-layer MLP to encode the prefix - self.embedding = torch.nn.Embedding(config["num_virtual_tokens"], config["token_dim"]) + self.embedding = torch.nn.Embedding(num_virtual_tokens, token_dim) self.trans = torch.nn.Sequential( - torch.nn.Linear(config["token_dim"], config["prompt_hidden_size"]), + torch.nn.Linear(token_dim, encoder_hidden_size), torch.nn.Tanh(), - torch.nn.Linear( - config["prompt_hidden_size"], - config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], - ), + torch.nn.Linear(encoder_hidden_size, num_layers * 2 * token_dim), ) else: - self.embedding = torch.nn.Embedding( - config["num_virtual_tokens"], - config["num_layers"] * 2 * config["num_transformer_submodules"] * config["token_dim"], - ) + self.embedding = torch.nn.Embedding(num_virtual_tokens, num_layers * 2 * token_dim) def forward(self, prefix: torch.Tensor): if self.prefix_projection: diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py index dca37c9..fdad421 100644 --- a/src/pet/tuners/prompt_tuning.py +++ b/src/pet/tuners/prompt_tuning.py @@ -1,25 +1,49 @@ import enum import math +from dataclasses import dataclass, field +from typing import Optional, Union import torch +from ..utils import PromptLearningConfig + class PromptTuningInit(str, enum.Enum): TEXT = "TEXT" RANDOM = "RANDOM" +@dataclass +class PromptTuningConfig(PromptLearningConfig): + prompt_tuning_init: Union[PromptTuningInit, str] = field( + default=PromptTuningInit.RANDOM, + metadata={"help": "How to initialize the prompt tuning parameters"}, + ) + prompt_tuning_init_text: Optional[str] = field( + default=None, + metadata={ + "help": "The text to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`" + }, + ) + tokenizer_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "The tokenizer to use for prompt tuning initialization. Only used if prompt_tuning_init is `TEXT`" + }, + ) + + class PromptEmbedding(torch.nn.Module): def __init__(self, config, word_embeddings): super().__init__() - total_virtual_tokens = config["num_virtual_tokens"] * config["num_transformer_submodules"] + total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"]) - if config["prompt_encoder_config"]["prompt_tuning_init"] == PromptTuningInit.TEXT: + if config.prompt_tuning_init == PromptTuningInit.TEXT: from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(config["prompt_encoder_config"]["tokenizer_name_or_path"]) - self.init_text = config["prompt_encoder_config"]["prompt_tuning_text"] + self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path) + self.init_text = config.prompt_tuning_init_text init_token_ids = self.tokenizer(self.init_text)["input_ids"] # Trim or iterate until num_text_tokens matches total_virtual_tokens num_text_tokens = len(init_token_ids) diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py new file mode 100644 index 0000000..2a32be1 --- /dev/null +++ b/src/pet/utils/__init__.py @@ -0,0 +1,5 @@ +# flake8: noqa +# There's no way to ignore "F401 '...' imported but unused" warnings in this +# module, but to preserve other warnings. So, don't check this module at all + +from .config import PETConfig, PETType, PromptLearningConfig, TaskType diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py new file mode 100644 index 0000000..6f66966 --- /dev/null +++ b/src/pet/utils/config.py @@ -0,0 +1,36 @@ +import enum +from dataclasses import dataclass, field +from typing import Optional, Union + + +class PETType(str, enum.Enum): + PROMPT_TUNING = "PROMPT_TUNING" + P_TUNING = "P_TUNING" + PREFIX_TUNING = "PREFIX_TUNING" + LORA = "LORA" + + +class TaskType(str, enum.Enum): + SEQ_CLS = "SEQ_CLS" + SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM" + CAUSAL_LM = "CAUSAL_LM" + + +@dataclass +class PETConfig: + """ + This is the configuration class to store the configuration of a :class:`~transform + """ + + pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"}) + task_type: Union[str, TaskType] = field(default=None, metadata={"help": "Task type"}) + inference_mode: bool = field(default=False, metadata={"help": "Whether to use inference mode"}) + + +@dataclass +class PromptLearningConfig(PETConfig): + num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"}) + token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"}) + num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"}) + num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"}) + num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"}) diff --git a/src/pet/utils/constants.py b/src/pet/utils/constants.py index 044a482..e832fa3 100644 --- a/src/pet/utils/constants.py +++ b/src/pet/utils/constants.py @@ -1 +1 @@ -# todo +# ToDo diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py new file mode 100644 index 0000000..45d7da1 --- /dev/null +++ b/src/pet/utils/other.py @@ -0,0 +1,14 @@ +import torch + + +def bloom_model_postprocess_past_key_value(past_key_values): + past_key_values = torch.cat(past_key_values) + total_layers, batch_size, num_attention_heads, num_virtual_tokens, head_dim = past_key_values.shape + keys = past_key_values[: total_layers // 2] + keys = keys.transpose(2, 3).reshape( + total_layers // 2, batch_size * num_attention_heads, head_dim, num_virtual_tokens + ) + values = past_key_values[total_layers // 2 :] + values = values.reshape(total_layers // 2, batch_size * num_attention_heads, num_virtual_tokens, head_dim) + + return tuple(zip(keys, values))