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

refactor

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-11-29 18:40:41 +05:30
parent 7c21cb456c
commit 91deee81e2
2 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pet/__init__.py
+1 -1
View File
@@ -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,
src/pet/task_mapping.py → src/pet/mapping.py
+9 -1
View File
@@ -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)