wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-11 12:30:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

adding detailed docs, refactor and fixes

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-12-01 14:34:09 +05:30
parent a21f24d9e1
commit e3d6568a19
11 changed files with 391 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
setup.py
+2 -1
View File
@@ -39,9 +39,10 @@ setup(
"packaging>=20.0",
"psutil",
"pyyaml",
"torch>=1.4.0",
"torch>=1.13.0",
"transformers",
"accelerate",
"loralib",
],
extras_require=extras,
classifiers=[
src/pet/mapping.py
+15
View File
@@ -34,6 +34,13 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
def get_pet_config(config_dict):
"""
Returns a PET config object from a dictionary.
Args:
config_dict (:obj:`Dict[str, Any]`):
"""
return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
@@ -93,6 +100,14 @@ def _prepare_lora_config(pet_config, model_config):
def get_pet_model(model, pet_config):
"""
Returns a PET model object from a model and a config.
Args:
model (:obj:`transformers.PreTrainedModel`):
pet_config (:obj:`transformers.PETConfig`):
"""
model_config = model.config.to_dict()
if pet_config.pet_type != PETType.LORA:
pet_config = _prepare_prompt_learning_config(pet_config, model_config)
src/pet/pet_model.py
+136 -5
View File
@@ -7,10 +7,30 @@ from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
from .utils import PETConfig, PETType, TaskType, shift_tokens_right
from .utils import PETConfig, PETType, TaskType, shift_tokens_right, _set_trainable
class PETModel(torch.nn.Module):
"""
Parameter Efficient Tuning Model. Base model encompassing various PET methods.
Args:
model (:obj:`PreTrainedModel`): The base transformer model used for PET.
pet_config (:obj:`PETConfig`): The configuration of the PET model.
Attributes:
base_model (:obj:`PreTrainedModel`): The base transformer model used for PET.
pet_config (:obj:`PETConfig`): The configuration of the PET model.
modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names to save when saving the model.
prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if `pet_config.pet_type != PETType.LORA`.
prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for PET if `pet_config.pet_type != PETType.LORA`.
transformer_backbone_name (:obj:`str`): The name of the transformer backbone in the base model
if `pet_config.pet_type != PETType.LORA`.
word_embeddings (:obj:`torch.nn.Embedding`): The word embeddings of the transformer backbone
in the base model if `pet_config.pet_type != PETType.LORA`.
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__()
self.pet_config = pet_config
@@ -54,6 +74,10 @@ class PETModel(torch.nn.Module):
).long()
def get_prompt_embedding_to_save(self):
"""
Returns the prompt embedding to save when saving the model.
Only applocable when `pet_config.pet_type != PETType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
@@ -61,6 +85,10 @@ class PETModel(torch.nn.Module):
return prompt_embeddings[0].detach().cpu()
def get_prompt(self, batch_size):
"""
Returns the virtual prompts to use for PET.
Only applocable when `pet_config.pet_type != PETType.LORA`.
"""
prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
@@ -92,6 +120,9 @@ class PETModel(torch.nn.Module):
return prompts
def print_trainable_parameters(self):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in self.named_parameters():
@@ -104,6 +135,41 @@ class PETModel(torch.nn.Module):
class PETModelForSequenceClassification(PETModel):
"""
PET model for sequence classification tasks.
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Attributes:
config (:obj:`PretrainedConfig`): The configuration object of the base model.
cls_layer_name (:obj:`str`): The name of the classification layer.
Example::
>>> from transformers import AutoModelForSequenceClassification
>>> from pet import PETModelForSequenceClassification, get_pet_config
>>> config = {
'pet_type': 'PREFIX_TUNING',
'task_type': 'SEQ_CLS',
'inference_mode': False,
'num_virtual_tokens': 20,
'token_dim': 768,
'num_transformer_submodules': 1,
'num_attention_heads': 12,
'num_layers': 12,
'encoder_hidden_size': 768,
'prefix_projection': False,
'postprocess_past_key_value_function': None
}
>>> pet_config = get_pet_config(config)
>>> model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased")
>>> pet_model = PETModelForSequenceClassification(model, pet_config)
>>> pet_model.print_trainable_parameters()
trainable params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
@@ -114,6 +180,9 @@ class PETModelForSequenceClassification(PETModel):
self.cls_layer_name = name
break
# to make sure classifier layer is trainable
_set_trainable(self.base_model)
def forward(
self,
input_ids=None,
@@ -160,7 +229,7 @@ class PETModelForSequenceClassification(PETModel):
)
if self.pet_config.pet_type == PETType.PREFIX_TUNING:
return self.prefix_tuning_forward(input_ids=input_ids, **kwargs)
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(
@@ -176,7 +245,7 @@ class PETModelForSequenceClassification(PETModel):
inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
def prefix_tuning_forward(
def _prefix_tuning_forward(
self,
input_ids=None,
attention_mask=None,
@@ -249,9 +318,40 @@ class PETModelForSequenceClassification(PETModel):
class PETModelForCausalLM(PETModel):
"""
PET model for Causal LM
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Example::
>>> from transformers import AutoModelForCausalLM
>>> from pet import PETModelForCausalLM, get_pet_config
>>> config = {
'pet_type': 'PREFIX_TUNING',
'task_type': 'CAUSAL_LM',
'inference_mode': False,
'num_virtual_tokens': 20,
'token_dim': 1280,
'num_transformer_submodules': 1,
'num_attention_heads': 20,
'num_layers': 36,
'encoder_hidden_size': 1280,
'prefix_projection': False,
'postprocess_past_key_value_function': None
}
>>> pet_config = get_pet_config(config)
>>> model = AutoModelForCausalLM.from_pretrained("gpt2-large")
>>> pet_model = PETModelForCausalLM(model, pet_config)
>>> pet_model.print_trainable_parameters()
trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
def forward(
self,
@@ -318,9 +418,40 @@ class PETModelForCausalLM(PETModel):
class PETModelForSeq2SeqLM(PETModel):
"""
PET model for Seq2Seq LM
Args:
model (:obj:`PreTrainedModel`): Base transformer model
pet_config (:obj:`PETConfig`): PET config.
Example::
>>> from transformers import AutoModelForSeq2SeqLM
>>> from pet import PETModelForSeq2SeqLM, get_pet_config
>>> config = {
'pet_type': 'LORA',
'task_type': 'SEQ_2_SEQ_LM',
'inference_mode': False,
'r': 8,
'target_modules': ['q', 'v'],
'lora_alpha': 32,
'lora_dropout': 0.1,
'merge_weights': False,
'fan_in_fan_out': False,
'enable_lora': None,
'bias': 'none'
}
>>> pet_config = get_pet_config(config)
>>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
>>> pet_model = PETModelForSeq2SeqLM(model, pet_config)
>>> pet_model.print_trainable_parameters()
trainable params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566
"""
def __init__(self, model, pet_config: PETConfig):
super().__init__(model, pet_config)
self.config = self.base_model.config
def forward(
self,
src/pet/tuners/lora.py
+54 -9
View File
@@ -13,6 +13,21 @@ from ..utils import PETConfig
@dataclass
class LoRAConfig(PETConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.LoRA`.
Args:
r: (:obj:`init`): LoRA attention dimension
target_modules (:obj: list of :obj: str): The names of the modules to apply LoRA to.
lora_alpha (:obj: float): The alpha parameter for LoRA scaling.
lora_dropout (:obj: float): The dropout probability for LoRA layers.
merge_weights (:obj: bool):
Whether to merge the weights of the LoRA layers with the base transformer model in `eval` mode.
fan_in_fan_out (:obj: bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
enable_lora (:obj: list of :obj: bool): Used with `lora.MergedLinear`.
bias (:obj: str): Bias type for LoRA. Can be 'none', 'all' or 'lora_only'
"""
r: int = field(default=8, metadata={"help": "LoRA attention dimension"})
target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"})
@@ -29,14 +44,44 @@ class LoRAConfig(PETConfig):
class LoRAModel(torch.nn.Module):
"""
Creates Low Rank Adapter (LoRA) model from a pretrained transformers model.
Args:
model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
config (:obj:`LoRAConfig`): The configuration of the LoRA model.
Returns:
:obj:`torch.nn.Module`: The LoRA model.
Example::
>>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig
>>> from pet import LoRAModel, LoRAConfig
>>> config = LoRAConfig(
pet_type="LORA",
task_type="SEQ_2_SEQ_LM",
r=8,
lora_alpha=32,
target_modules=["q", "v"],
lora_dropout=0.01,
)
>>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
>>> lora_model = LoRAModel(config, model)
Attributes:
model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
config (:obj:`LoRAConfig`): The configuration of the LoRA model.
"""
def __init__(self, config, model):
super().__init__()
self.config = config
self.model = model
self.find_and_replace()
self._find_and_replace()
mark_only_lora_as_trainable(self.model, self.config.bias)
def find_and_replace(self):
def _find_and_replace(self):
kwargs = {
"r": self.config.r,
"lora_alpha": self.config.lora_alpha,
@@ -47,23 +92,23 @@ class LoRAModel(torch.nn.Module):
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_modules):
parent, target, target_name = self.get_submodules(key)
# print(parent, target, target_name)
parent, target, target_name = self._get_submodules(key)
bias = target.bias is not None
if isinstance(target, torch.nn.Linear):
new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
new_module = lora.Linear(target.in_features, target.out_features, bias=bias, **kwargs)
elif isinstance(target, Conv1D):
kwargs.update({"enable_lora": self.config.enable_lora})
in_features, out_features = target.weight.shape
new_module = lora.MergedLinear(in_features, out_features, **kwargs)
self.replace_module(parent, target_name, new_module, target)
new_module = lora.MergedLinear(in_features, out_features, bias=bias, **kwargs)
self._replace_module(parent, target_name, new_module, target)
def get_submodules(self, key):
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, target_name
def replace_module(self, parent_module, child_name, new_module, old_module):
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
if old_module.bias is not None:
src/pet/tuners/p_tuning.py
+52 -3
View File
@@ -14,21 +14,33 @@ class PromptEncoderReparameterizationType(str, enum.Enum):
@dataclass
class PromptEncoderConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PromptEncoder`.
Args:
encoder_reparameterization_type
(:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
The type of reparameterization to use.
encoder_hidden_size (:obj:`int`): The hidden size of the prompt encoder.
encoder_num_layers (:obj:`int`): The number of layers of the prompt encoder.
encoder_dropout (:obj:`float`): The dropout probability of the prompt encoder.
"""
encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field(
default=PromptEncoderReparameterizationType.MLP,
metadata={"help": "How to reparameterize the prompt encoder"},
)
encoder_hidden_size: int = field(
default=None,
metadata={"help": "The hidden size of the prompt encoder reparameterization"},
metadata={"help": "The hidden size of the prompt encoder"},
)
encoder_num_layers: int = field(
default=2,
metadata={"help": "The number of layers of the prompt encoder reparameterization"},
metadata={"help": "The number of layers of the prompt encoder"},
)
encoder_dropout: float = field(
default=0.0,
metadata={"help": "The dropout of the prompt encoder reparameterization"},
metadata={"help": "The dropout of the prompt encoder"},
)
@@ -37,6 +49,43 @@ class PromptEncoderConfig(PromptLearningConfig):
class PromptEncoder(torch.nn.Module):
"""
The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
Args:
config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder.
Example::
>>> from pet import PromptEncoder, PromptEncoderConfig
>>> config = PromptEncoderConfig(
pet_type="P_TUNING",
task_type="SEQ_2_SEQ_LM",
num_virtual_tokens=20,
token_dim=768,
num_transformer_submodules=1,
num_attention_heads=12,
num_layers=12,
encoder_reparameterization_type="MLP",
encoder_hidden_size=768
)
>>> prompt_encoder = PromptEncoder(config)
Attributes:
embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder.
mlp_head (:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`.
lstm_head (:class:`~torch.nn.LSTM`):
The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`.
token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
input_size (:obj:`int`): The input size of the prompt encoder.
output_size (:obj:`int`): The output size of the prompt encoder.
hidden_size (:obj:`int`): The hidden size of the prompt encoder.
total_virtual_tokens (:obj:`int`): The total number of virtual tokens of the prompt encoder.
encoder_type (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
The encoder type of the prompt encoder.
Input shape: (batch_size, total_virtual_tokens)
Output shape: (batch_size, total_virtual_tokens, token_dim)
"""
def __init__(self, config):
src/pet/tuners/prefix_tuning.py
+34
View File
@@ -8,6 +8,15 @@ from ..utils import PromptLearningConfig
@dataclass
class PrefixTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PrefixEncoder`.
Args:
encoder_hidden_size (:obj: int): The hidden size of the prompt encoder.
prefix_projection (:obj: bool): Whether to project the prefix embeddings.
postprocess_past_key_value_function (:obj: Optional[Callable]): The function to postprocess the past key value.
"""
encoder_hidden_size: int = field(
default=None,
metadata={"help": "The hidden size of the encoder"},
@@ -28,6 +37,31 @@ class PrefixEncoder(torch.nn.Module):
r"""
The torch.nn model to encode the prefix
Args:
config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder.
Example::
>>> from pet import PrefixEncoder, PrefixTuningConfig
>>> config = PrefixTuningConfig(
pet_type="PREFIX_TUNING",
task_type="SEQ_2_SEQ_LM",
num_virtual_tokens=20,
token_dim=768,
num_transformer_submodules=1,
num_attention_heads=12,
num_layers=12,
encoder_hidden_size=768
)
>>> prefix_encoder = PrefixEncoder(config)
Attributes:
embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder.
trans (:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings
if :obj:`prefix_projection` is :obj:`True`.
prefix_projection (:obj:`bool`): Whether to project the prefix embeddings.
Input shape: (batch_size, num_virtual_tokens)
Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden)
src/pet/tuners/prompt_tuning.py
+48 -3
View File
@@ -15,6 +15,17 @@ class PromptTuningInit(str, enum.Enum):
@dataclass
class PromptTuningConfig(PromptLearningConfig):
"""
This is the configuration class to store the configuration of a :class:`~pet.PromptEmbedding`.
Args:
prompt_tuning_init (:obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding.
prompt_tuning_init_text (:obj: Optional[:obj:`str`]): The text to initialize the prompt embedding.
Only used if `prompt_tuning_init` is `TEXT`
tokenizer_name_or_path (:obj: Optional[:obj:`str`]): The name or path of the tokenizer.
Only used if `prompt_tuning_init` is `TEXT`
"""
prompt_tuning_init: Union[PromptTuningInit, str] = field(
default=PromptTuningInit.RANDOM,
metadata={"help": "How to initialize the prompt tuning parameters"},
@@ -34,6 +45,40 @@ class PromptTuningConfig(PromptLearningConfig):
class PromptEmbedding(torch.nn.Module):
"""
The model to encode virtual tokens into prompt embeddings.
Args:
config (:class:`PromptTuningConfig`): The configuration of the prompt embedding.
word_embeddings (:obj:`torch.nn.Module`): The word embeddings of the base transformer model.
Attributes:
embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prompt embedding.
Example::
>>> from pet import PromptEmbedding, PromptTuningConfig
>>> config = PromptTuningConfig(
pet_type="PROMPT_TUNING",
task_type="SEQ_2_SEQ_LM",
num_virtual_tokens=20,
token_dim=768,
num_transformer_submodules=1,
num_attention_heads=12,
num_layers=12,
prompt_tuning_init="TEXT",
prompt_tuning_init_text="Predict if sentiment of this review is positive, negative or neutral",
tokenizer_name_or_path="t5-base",
)
>>> # t5_model.shared is the word embeddings of the base model
>>> prompt_embedding = PromptEmbedding(config, t5_model.shared)
Input Shape: (batch_size, total_virtual_tokens)
Output Shape: (batch_size, total_virtual_tokens, token_dim)
"""
def __init__(self, config, word_embeddings):
super().__init__()
@@ -42,9 +87,9 @@ class PromptEmbedding(torch.nn.Module):
if config.prompt_tuning_init == PromptTuningInit.TEXT:
from transformers import AutoTokenizer
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"]
tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
init_text = config.prompt_tuning_init_text
init_token_ids = tokenizer(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:
src/pet/utils/__init__.py
+1 -1
View File
@@ -3,5 +3,5 @@
# module, but to preserve other warnings. So, don't check this module at all
from .config import PETConfig, PETType, PromptLearningConfig, TaskType
from .other import bloom_model_postprocess_past_key_value, shift_tokens_right
from .other import bloom_model_postprocess_past_key_value, shift_tokens_right, _set_trainable
from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict
src/pet/utils/config.py
+20 -2
View File
@@ -19,7 +19,12 @@ class TaskType(str, enum.Enum):
@dataclass
class PETConfig:
"""
This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
This is the base configuration class to store the configuration of a :class:`~pet.PETModel`.
Args:
pet_type (:obj:Union[:class:`~pet.utils.config.PETType`, :obj:`str`]): The type of PET method to use.
task_type (:obj:Union[:class:`~pet.utils.config.TaskType`, :obj:`str`]): The type of task to perform.
inference_mode (:obj:`bool`, defaults to :obj:`False`): Whether to use the PET model in inference mode.
"""
pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
@@ -29,8 +34,21 @@ class PETConfig:
@dataclass
class PromptLearningConfig(PETConfig):
"""
This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`, :class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`].
Args:
num_virtual_tokens (:obj:`int`): The number of virtual tokens to use.
token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
num_transformer_submodules (:obj:`int`): The number of transformer submodules in the base transformer model.
num_attention_heads (:obj:`int`): The number of attention heads in the base transformer model.
num_layers (:obj:`int`): The number of layers in the base transformer model.
"""
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"})
token_dim: int = field(
default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
)
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/other.py
+14
View File
@@ -19,6 +19,11 @@ def bloom_model_postprocess_past_key_value(past_key_values):
def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
"""
Shift input ids one token to the right.
Args:
input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): input ids
pad_token_id (:obj:`int`): The id of the `padding` token.
decoder_start_token_id (:obj:`int`): The id of the `start` token.
"""
shifted_input_ids = input_ids.new_zeros(input_ids.shape)
shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
@@ -30,3 +35,12 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
return shifted_input_ids
def _set_trainable(model):
if model.modules_to_save is not None:
for name, param in model.named_parameters():
if any(module_name in name for module_name in model.modules_to_save):
param.requires_grad = True
else:
param.requires_grad = False
src/pet/utils/save_and_load.py
+15
View File
@@ -4,6 +4,13 @@ from .config import PETType
def get_pet_model_state_dict(model):
"""
Get the state dict of the PET model.
Args:
model (:obj:`PETModel`): The PET model.
"""
if model.pet_config.pet_type == PETType.LORA:
return lora_state_dict(model)
else:
@@ -19,6 +26,14 @@ def get_pet_model_state_dict(model):
def set_pet_model_state_dict(model, pet_model_state_dict):
"""
Set the state dict of the PET model.
Args:
model (:obj:`PETModel`): The PET model.
pet_model_state_dict (:obj:`dict`): The state dict of the PET model.
"""
model.load_state_dict(pet_model_state_dict, strict=False)
if model.pet_config.pet_type != PETType.LORA:
model.prompt_encoder.embedding.load_state_dict(