lot of refactoring

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      | ✅  | ✅  | ✅  |   |
 
 
 

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

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)

src/pet/prompt_learning_legacy.py

@@ -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()

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,
+}

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

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()

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,

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:

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)

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

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"})

src/pet/utils/constants.py

@@ -1 +1 @@
-# todo
+# ToDo

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))