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

refactor

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-11-30 18:58:28 +05:30
parent 3626d4cf76
commit 207eb07f03
1 changed files with 27 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pet/mapping.py
+27 -27
View File
@@ -37,38 +37,38 @@ def get_pet_config(config_dict):
return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
def _prepare_prompt_learning_config(pet_config, config):
def _prepare_prompt_learning_config(pet_config, model_config):
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"]
if "num_hidden_layers" in model_config:
num_layers = model_config["num_hidden_layers"]
elif "num_layers" in model_config:
num_layers = model_config["num_layers"]
elif "n_layer" in model_config:
num_layers = model_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"]
if "hidden_size" in model_config:
token_dim = model_config["hidden_size"]
elif "n_embd" in model_config:
token_dim = model_config["n_embd"]
elif "d_model" in model_config:
token_dim = model_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"]
if "num_attention_heads" in model_config:
num_attention_heads = model_config["num_attention_heads"]
elif "n_head" in model_config:
num_attention_heads = model_config["n_head"]
elif "num_heads" in model_config:
num_attention_heads = model_config["num_heads"]
elif "encoder_attention_heads" in model_config:
num_attention_heads = model_config["encoder_attention_heads"]
else:
raise ValueError("Please specify `num_attention_heads` in `pet_config`")
pet_config.num_attention_heads = num_attention_heads
@@ -79,11 +79,11 @@ def _prepare_prompt_learning_config(pet_config, config):
return pet_config
def _prepare_lora_config(pet_config, config):
def _prepare_lora_config(pet_config, model_config):
if pet_config.target_modules is None:
if config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
if model_config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
raise ValueError("Please specify `target_modules` in `pet_config`")
pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[config.model_type]
pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config.model_type]
if len(pet_config.target_modules) == 1:
pet_config.fan_in_fan_out = True
pet_config.enable_lora = [True, False, True]
@@ -93,10 +93,10 @@ def _prepare_lora_config(pet_config, config):
def get_pet_model(model, pet_config):
config = model.config.to_dict()
model_config = model.config.to_dict()
if pet_config.pet_type != PETType.LORA:
pet_config = _prepare_prompt_learning_config(pet_config, config)
pet_config = _prepare_prompt_learning_config(pet_config, model_config)
else:
pet_config = _prepare_lora_config(pet_config, config)
pet_config = _prepare_lora_config(pet_config, model_config)
return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config)