diff --git a/src/pet/__init__.py b/src/pet/__init__.py
index 51ab10c..43b71b2 100644
--- a/src/pet/__init__.py
+++ b/src/pet/__init__.py
@@ -5,7 +5,7 @@
 __version__ = "0.1.0.dev0"
 
 from .pet_model import PETModel, PETModelForCausalLM, PETModelForSeq2SeqLM, PETModelForSequenceClassification
-from .task_mapping import MODEL_TYPE_TO_PROMPT_MODEL_MAPPING
+from .mapping import MODEL_TYPE_TO_PET_MODEL_MAPPING, PET_TYPE_TO_CONFIG_MAPPING, get_pet_config, get_pet_model
 from .tuners import (
     PrefixEncoder,
     PrefixTuningConfig,
diff --git a/src/pet/task_mapping.py b/src/pet/mapping.py
similarity index 65%
rename from src/pet/task_mapping.py
rename to src/pet/mapping.py
index 00e7415..39f7b89 100644
--- a/src/pet/task_mapping.py
+++ b/src/pet/mapping.py
@@ -3,7 +3,7 @@ from .tuners import PrefixTuningConfig, PromptEncoderConfig, PromptTuningConfig
 from .utils import PETConfig
 
 
-MODEL_TYPE_TO_PROMPT_MODEL_MAPPING = {
+MODEL_TYPE_TO_PET_MODEL_MAPPING = {
     "SEQ_CLS": PETModelForSequenceClassification,
     "SEQ_2_SEQ_LM": PETModelForSeq2SeqLM,
     "CAUSAL_LM": PETModelForCausalLM,
@@ -15,3 +15,11 @@ PET_TYPE_TO_CONFIG_MAPPING = {
     "P_TUNING": PromptEncoderConfig,
     "LORA": PETConfig,
 }
+
+
+def get_pet_config(config_dict):
+    return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
+
+
+def get_pet_model(model, pet_config):
+    return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config)