Multi_class_classification errors when fine-tuning via TrainerAPI
I’m attempting to do a multiclass finetuning of distilbert using the following code:import torchfrom transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer, DataCollatorWithPaddingfrom datasets import load_datasetimport evaluatefrom tqdm.notebook import tqdmfrom datasets import load_datasetimport numpy as npimport osdataset = load_dataset(“csv”, data_files=“datafile.csv”,keep_default_na=False)dataset[“train”] = dataset[“train”].rename_column(“label”, “labels”)dataset = dataset[“train”].train_test_split(test_size=0.3)dataset[“train”] = dataset[“train”].shuffle()dataset[“test”] = dataset[“test”].shuffle()def preprocess_function(examples):return tokenizer(examples[“text”], truncation=True)model_id = “distilbert-base-uncased”tokenizer = AutoTokenizer.from_pretrained(model_id, max_length = 512, model_max_length = 512)tokenized_dataset = dataset.map(preprocess_function, batched=True)data_collator = DataCollatorWithPadding(tokenizer=tokenizer)accuracy = evaluate.load(“accuracy”)def compute_metrics(eval_pred):predictions, labels = eval_predpredictions = np.argmax(predictions, axis=1)return accuracy.compute(predictions=predictions, references=labels)id2label = {0: “casual”, 1: “possibly_needs_caution”, 2: “probably_needs_caution”, 3: “needs_caution”, 4: “needs_intervention”}label2id = {“casual”: 0, “possibly_needs_caution”: 1, “probably_needs_caution”: 2, “needs_caution”: 3, “needs_intervention”: 4}model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=5, id2label=id2label, label2id=label2id,problem_type="multi_class_classification")tokenized_dataset[“train”] = tokenized_dataset[“train”].remove_columns([‘text’,‘Unnamed: 0’])tokenized_dataset[“test”] = tokenized_dataset[“test”].remove_columns([‘text’,‘Unnamed: 0’])training_args = TrainingArguments(output_dir="safety_instruct_detect",

learning_rate=4.5e-5,

per_device_train_batch_size=8,

per_device_eval_batch_size=8,

#fp16 = True,

warmup_steps = 1000, 

num_train_epochs=2,

evaluation_strategy="epoch",

save_strategy="epoch",

load_best_model_at_end=True,

push_to_hub=False,)trainer = Trainer(model=model,

args=training_args,

train_dataset=tokenized_dataset["train"],

eval_dataset=tokenized_dataset["test"],

tokenizer=tokenizer,

data_collator=data_collator,


compute_metrics=compute_metrics,)trainer.train()I get the following error:ValueError: The model did not return a loss from the inputs, only the following keys: logits. For reference, the inputs it received are labels,input_ids,attention_mask.Here are a few things i’ve tried:Using BertForSequenceClassification rather than AutoModelForSequenceClassificationMaking the labels in the dataset a float rather than intUsing both “label” and “labels” as column names a few different times