Peft Prompt Tuning - ValueError: `create_and_replace` does not support prompt learning and adaption prompt yet

I am trying to perform prompt tuning for classification task as shown below running into this error:

ValueError                                Traceback (most recent call last)
File <command-1698098128188718>, line 13
      9 trainer.train()
     10 trainer.save_model(model_output_dir)
     11 pipe = pipeline(
     12     task="sequence-classification",
---> 13     model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), # this needs fixing
     14     batch_size=8,
     15     tokenizer=tokenizer,
     16 )
     17 pipe.save_pretrained(pipeline_output_dir)
     18 mlflow.transformers.log_model(
     19     transformers_model=pipe,
     20     artifact_path=model_artifact_path,
     21     input_example="Hi there!",
     22 )

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-753e0642-e159-43ba-a3e9-89dbaed5c6ed/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:566, in _BaseAutoModelClass.from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
    564 elif type(config) in cls._model_mapping.keys():
    565     model_class = _get_model_class(config, cls._model_mapping)
--> 566     return model_class.from_pretrained(
    567         pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
    568     )
    569 raise ValueError(
    570     f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
    571     f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
    572 )

File /databricks/python_shell/dbruntime/huggingface_patches/transformers.py:21, in _create_patch_function.<locals>.patched_from_pretrained(cls, *args, **kwargs)
     19 call_succeeded = False
     20 try:
---> 21     model = original_method.__func__(cls, *args, **kwargs)
     22     call_succeeded = True
     23     return model

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-753e0642-e159-43ba-a3e9-89dbaed5c6ed/lib/python3.10/site-packages/transformers/modeling_utils.py:3787, in PreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, use_safetensors, *model_args, **kwargs)
   3784     model = quantizer.post_init_model(model)
   3786 if _adapter_model_path is not None:
-> 3787     model.load_adapter(
   3788         _adapter_model_path,
   3789         adapter_name=adapter_name,
   3790         token=token,
   3791         adapter_kwargs=adapter_kwargs,
   3792     )
   3794 if output_loading_info:
   3795     if loading_info is None:

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-753e0642-e159-43ba-a3e9-89dbaed5c6ed/lib/python3.10/site-packages/transformers/integrations/peft.py:187, in PeftAdapterMixin.load_adapter(self, peft_model_id, adapter_name, revision, token, device_map, max_memory, offload_folder, offload_index, peft_config, adapter_state_dict, adapter_kwargs)
    180     peft_config = PeftConfig.from_pretrained(
    181         peft_model_id,
    182         token=token,
    183         **adapter_kwargs,
    184     )
    186 # Create and add fresh new adapters into the model.
--> 187 inject_adapter_in_model(peft_config, self, adapter_name)
    189 if not self._hf_peft_config_loaded:
    190     self._hf_peft_config_loaded = True

File /local_disk0/.ephemeral_nfs/envs/pythonEnv-753e0642-e159-43ba-a3e9-89dbaed5c6ed/lib/python3.10/site-packages/peft/mapping.py:153, in inject_adapter_in_model(peft_config, model, adapter_name)
    139 r"""
    140 A simple API to create and inject adapter in-place into a model. Currently the API does not support prompt learning
    141 methods and adaption prompt. Make sure to have the correct `target_names` set in the `peft_config` object. The API
   (...)
    150         The name of the adapter to be injected, if not provided, the default adapter name is used ("default").
    151 """
    152 if peft_config.is_prompt_learning or peft_config.is_adaption_prompt:
--> 153     raise ValueError("`create_and_replace` does not support prompt learning and adaption prompt yet.")
    155 if peft_config.peft_type not in PEFT_TYPE_TO_TUNER_MAPPING.keys():
    156     raise ValueError(
    157         f"`inject_adapter_in_model` does not support {peft_config.peft_type} yet. Please use `get_peft_model`."
    158     )

ValueError: `create_and_replace` does not support prompt learning and adaption prompt yet.
from transformers import pipeline, TextClassificationPipeline
import mlflow
from transformers import (
    Trainer,
    DataCollatorForLanguageModeling,
    DataCollatorWithPadding,
)

from peft import (
    get_peft_config,
    get_peft_model,
    PromptTuningInit,
    PromptTuningConfig,
    PromptEmbedding,
    TaskType,
    PeftType,
    AutoPeftModelForSequenceClassification,
    AutoPeftModelForSeq2SeqLM,
    AutoPeftModel,
)

from transformers import AutoModelForSequenceClassification
from transformers import TrainingArguments
import evaluate
from transformers import AutoTokenizer

model_name = "distilbert-base-uncased" 

tokenizer = AutoTokenizer.from_pretrained(model_name)

def tokenize_function(examples):
    return tokenizer(
        examples["input"], padding="max_length", truncation=True, max_length=128
    )


train_tokenized = train_dataset.map(tokenize_function, batched=True).remove_columns(
    ["input"]
)
test_tokenized = test_dataset.map(tokenize_function, batched=True).remove_columns(
    ["input"]
)

train_data = train_tokenized.shuffle(seed=42)
test_data = test_tokenized.shuffle(seed=42)

metric = evaluate.load("accuracy")  # we need to make sure data is balanced

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

training_args = TrainingArguments(
    output_dir=output_directory,  # Where the model predictions and checkpoints will be written
    evaluation_strategy="epoch",
)

foundation_model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2, label2id=label2id, id2label=id2label
)
config = PromptTuningConfig(
    peft_type="PROMPT_TUNING",
    task_type="SEQ_CLS",  # SEQ_2_SEQ_LM", # "SEQ_CLS",
    num_virtual_tokens=3,
    token_dim=768,
    num_transformer_submodules=1,
    num_attention_heads=12,
    num_layers=12,
    prompt_tuning_init=PromptTuningInit.TEXT,
    prompt_tuning_init_text="Classify if the input text of File name and Field names is supply chain or product related table:",
    tokenizer_name_or_path=base_model,
)
peft_model = get_peft_model(foundation_model, config)

trainer = Trainer(
    model=peft_model,  # We pass in the PEFT version of the foundation model
    args=training_args,
    train_dataset=train_data,
    eval_dataset=test_data,
    compute_metrics=compute_metrics,
    data_collator=DataCollatorWithPadding(
        tokenizer
    ),  # mlm=False indicates not to use masked language modeling , DataCollatorForLanguageModeling(tokenizer, mlm=False)
)

pipeline_artifact_name = "pipeline"


class TextClassificationPipelineModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        device = 0 if torch.cuda.is_available() else -1
        self.pipeline = pipeline(
            "text-classification",
            context.artifacts[pipeline_artifact_name],
            device=device,
        )

    def predict(self, context, model_input):
        texts = model_input[model_input.columns[0]].to_list()
        pipe = tqdm(
            self.pipeline(texts, truncation=True, batch_size=8),
            total=len(texts),
            miniters=10,
        )
        labels = [prediction["label"] for prediction in pipe]
        return pd.Series(labels)

with mlflow.start_run() as run:
    trainer.train()
    trainer.save_model(model_output_dir)
    pipe = pipeline(
        task="sequence-classification",
        model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), # this needs fixing
        batch_size=8,
        tokenizer=tokenizer,
    )
    pipe.save_pretrained(pipeline_output_dir)
    mlflow.transformers.log_model(
        transformers_model=pipe,
        artifact_path=model_artifact_path,
        input_example="Hi there!",
    )