PEFT prompt tuning for SEQ_CLS with BERT causes unexpected keyword argument 'label'

Beginners
1

Hi everyone,

I would like to use PEFT’s prompt tuning on a BERT model for sequence classification (I just want to practice a bit of PEFT, not seeking high performances). I ran into a common problem about unexpected keyword argument “label” when I try to train my PEFT model.

It is worth mentioning that my dataset and my base-model are working in a LoRA framework.

I created a PeftConfig this way:

generation_config = PromptTuningConfig(
    task_type=TaskType.SEQ_CLS,
    prompt_tuning_init="TEXT",
    prompt_tuning_init_text="Classiy between positive sentiment or negative sentiments",
    num_virtual_tokens=16, 
    tokenizer_name_or_path="google-bert/bert-base-cased"
)

Then I create my PEFT model using

prompt_model = get_peft_model(
    BertForSequenceClassification.from_pretrained("google-bert/bert-base-cased", num_labels=2),
    generation_config
  )

My dataset has been obtained from a tokenized Pandas dataFrame and looks like

DatasetDict({
    train: Dataset({
        features: ['text', 'label', 'input_ids', 'token_type_ids', 'attention_mask'],
        num_rows: 5000
    })
    test: Dataset({
        features: ['text', 'label', 'input_ids', 'token_type_ids', 'attention_mask'],
        num_rows: 20000
    })
})

I finally plug my new prompt_model into a a Trainer:

trainer = Trainer(
    model=prompt_model,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["test"],
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False) 
)

Calling trainer.train() provokes the (common) error TypeError: BertForSequenceClassification.forward() got an unexpected keyword argument 'label'.

According to this post, I used the correct BERT variant BertForSequenceClassification that should handle the label feature.

Please tell me what am I doing wrong ? Thanks :slight_smile: