Do I have to write custom AutoModel transformers class in case "TypeError: NVEmbedModel.forward() got an unexpected keyword argument 'inputs_embeds'"

I am trying to finetune the nvidia/NV-Embed-v2 model from hugging face using lora from peft library. I am facing the "TypeError: NVEmbedModel.forward() got an unexpected keyword argument 'inputs_embeds'".

I checked the files of the model and it does expect the inputs_embeds keyword argument.

I am using FEATURE_EXTRACTION task type for peft config, and using AutoModel to load the model with 8bit quantization.

Code:

from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, TaskType
import pandas as pd
from datasets import Dataset

# load the base model
model_name = "nvidia/NV-Embed-v2"
base_model = AutoModel.from_pretrained(model_name, trust_remote_code=True, load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# define lora config
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type=TaskType.FEATURE_EXTRACTION
)

peft_model = get_peft_model(base_model, lora_config)

# load the data
df = pd.DataFrame([['tomato', 'tomato ketchup heinz', 1], ['tomato', '500gm salted chips', 0], ['tomato', 'tomato 500gm', 1], ['strawberry', 'strawberry ripe 200gm', 1]], columns=['search_term', 'product_string', 'score'])
data = Dataset.from_pandas(df)
data = data.train_test_split(test_size=0.25)

train_dataset = data["train"]
test_dataset = data["test"]


def preprocess_function(examples):
    return tokenizer(examples["product_string"], examples["search_term"], truncation=True, padding="max_length", max_length=128)

train_dataset_final = train_dataset.map(preprocess_function, batched=True)

train_dataset_final.set_format(
    type="torch", columns=["input_ids", "attention_mask", "score"]
)

# Train the model
training_args = TrainingArguments(
    output_dir='lora-nvembedv2',
    auto_find_batch_size=True,
    learning_rate= 3e-2,
    num_train_epochs=1
)

trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=train_dataset_final,
)

trainer.train()

This is similar to this issue. The reason is that TaskType.FEATURE_EXTRACTION expects an additional input_embeds input which this model does not provide (neither does CLIP, for example).

Remove the TaskType.FEATURE_EXTRACTION, basically setting task_type=None so that you receive a plain PeftModel instance instead of a PeftModelForFeatureExtraction and this problem should vanish.