fix prefix tuning config to remove function field as it cannot be converted to json

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,

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:

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

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

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):
     """