get_pet_model fn

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)

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(

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(