diff --git a/src/peft/__init__.py b/src/peft/__init__.py index b76dad8..beb9423 100644 --- a/src/peft/__init__.py +++ b/src/peft/__init__.py @@ -40,6 +40,7 @@ from .tuners import ( PromptTuningInit, ) from .utils import ( + TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, PeftConfig, PeftType, PromptLearningConfig, diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index a39b1c2..de0853e 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -27,6 +27,7 @@ from huggingface_hub import hf_hub_download from .tuners import LoraModel, PrefixEncoder, PromptEmbedding, PromptEncoder from .utils import ( + TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, WEIGHTS_NAME, PeftConfig, PeftType, @@ -218,8 +219,8 @@ class PeftModel(PushToHubMixin, torch.nn.Module): past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split( self.peft_config.num_transformer_submodules * 2 ) - if self.peft_config.postprocess_past_key_value_function is not None: - post_process_fn = self.peft_config.postprocess_past_key_value_function + if TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING.get(self.config.model_type, None) is not None: + post_process_fn = TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING[self.config.model_type] past_key_values = post_process_fn(past_key_values) return past_key_values else: diff --git a/src/peft/tuners/prefix_tuning.py b/src/peft/tuners/prefix_tuning.py index d925c49..fcb207c 100644 --- a/src/peft/tuners/prefix_tuning.py +++ b/src/peft/tuners/prefix_tuning.py @@ -15,7 +15,6 @@ from dataclasses import dataclass, field -from typing import Callable, Optional import torch @@ -30,7 +29,6 @@ class PrefixTuningConfig(PromptLearningConfig): Args: encoder_hidden_size (`int`): The hidden size of the prompt encoder. prefix_projection (`bool`): Whether to project the prefix embeddings. - postprocess_past_key_value_function (`Callable`, *optional*): The function to postprocess the past key value. """ encoder_hidden_size: int = field( @@ -41,10 +39,6 @@ class PrefixTuningConfig(PromptLearningConfig): default=False, metadata={"help": "Whether to project the prefix tokens"}, ) - postprocess_past_key_value_function: Optional[Callable] = field( - default=None, - metadata={"help": "The function to postprocess the past key value"}, - ) def __post_init__(self): self.peft_type = PeftType.PREFIX_TUNING diff --git a/src/peft/utils/__init__.py b/src/peft/utils/__init__.py index c418d3d..b451fb4 100644 --- a/src/peft/utils/__init__.py +++ b/src/peft/utils/__init__.py @@ -19,5 +19,11 @@ from .adapters_utils import CONFIG_NAME, WEIGHTS_NAME from .config import PeftConfig, PeftType, PromptLearningConfig, TaskType -from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right, transpose +from .other import ( + TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING, + _set_trainable, + bloom_model_postprocess_past_key_value, + shift_tokens_right, + transpose, +) from .save_and_load import get_peft_model_state_dict, peft_model_load_and_dispatch, set_peft_model_state_dict diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py index 14ab90e..3f8627e 100644 --- a/src/peft/utils/other.py +++ b/src/peft/utils/other.py @@ -30,6 +30,11 @@ def bloom_model_postprocess_past_key_value(past_key_values): return tuple(zip(keys, values)) +TRANSFORMERS_MODELS_TO_PREFIX_TUNING_POSTPROCESS_MAPPING = { + "bloom": bloom_model_postprocess_past_key_value, +} + + # copied from transformers.models.bart.modeling_bart def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """