diff --git a/src/pet/mapping.py b/src/pet/mapping.py index 39f7b89..4c82bea 100644 --- a/src/pet/mapping.py +++ b/src/pet/mapping.py @@ -22,4 +22,43 @@ def get_pet_config(config_dict): def get_pet_model(model, pet_config): + config = model.config.to_dict() + if pet_config.num_layers is None: + if "num_hidden_layers" in config: + num_layers = config["num_hidden_layers"] + elif "num_layers" in config: + num_layers = config["num_layers"] + elif "n_layer" in config: + num_layers = config["n_layer"] + else: + raise ValueError("Please specify `num_layers` in `pet_config`") + pet_config.num_layers = num_layers + + if pet_config.token_dim is None: + 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 `pet_config`") + pet_config.token_dim = token_dim + + if pet_config.num_attention_heads is None: + 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"] + elif "encoder_attention_heads" in config: + num_attention_heads = config["encoder_attention_heads"] + else: + raise ValueError("Please specify `num_attention_heads` in `pet_config`") + pet_config.num_attention_heads = num_attention_heads + + if pet_config.encoder_hidden_size is None: + pet_config.encoder_hidden_size = token_dim + return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config) diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py index 8f8e030..bb5a5a8 100644 --- a/src/pet/tuners/p_tuning.py +++ b/src/pet/tuners/p_tuning.py @@ -19,7 +19,7 @@ class PromptEncoderConfig(PromptLearningConfig): metadata={"help": "How to reparameterize the prompt encoder"}, ) encoder_hidden_size: int = field( - default=256, + default=None, metadata={"help": "The hidden size of the prompt encoder reparameterization"}, ) encoder_num_layers: int = field( diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py index 049138b..d709baf 100644 --- a/src/pet/tuners/prefix_tuning.py +++ b/src/pet/tuners/prefix_tuning.py @@ -9,7 +9,7 @@ from ..utils import PromptLearningConfig @dataclass class PrefixTuningConfig(PromptLearningConfig): encoder_hidden_size: int = field( - default=256, + default=None, metadata={"help": "The hidden size of the encoder"}, ) prefix_projection: bool = field(