KeyError: 'test' when trying to divide a custom dataset into train and test for fine-tuning

🤗Transformers
1

So, I am trying to fine-tune a Protbert Transformers with a .fasta file I have converted into a CSV. I have been following this tutorial: (Fine-tune a pretrained model).

Everything goes fine until I get to the training section. I have uploaded the CSV with the sequences but when I try to create small datasets for training and testing as they say in the tutorial I get: KeyError: 'test'. It is strange because it doesn’t say anything else. Therefore, I can’t continue with the tutorial to fine-tune the model.

The code I have for now:

dataset = load_dataset('csv', data_files='/content/sequences.csv')

tokenizer = BertTokenizer.from_pretrained("Rostlab/prot_bert", do_lower_case=False )

def tokenize_function(samples):
    return tokenizer(samples["Protein Sequence"], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000)) #Here is where I get the error.

#I can't continue with this part because of the error.
training_args = TrainingArguments(output_dir="/content/drive/MyDrive/Colab Notebooks/test_trainer", evaluation_strategy="epoch")

metric = load_metric("accuracy")

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

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_dataset,
    eval_dataset=small_eval_dataset,
    compute_metrics=compute_metrics,
)

trainer.train()

Does anyone know why this happens? How should I follow the tutorial?