微调Protbert划分数据集时出现KeyError: 'test'问题求助

问题根源 #

用load_dataset('csv', data_files='/content/sequences.csv')加载单个CSV文件时，🤗 Datasets库默认只会生成名为train的数据集拆分，不会自动创建test拆分，直接访问tokenized_datasets["test"]自然会触发键错误。

两种可行解决方案 #

方案1：加载时直接拆分数据集 #

调用load_dataset时，利用train_test_split参数一次性完成数据集拆分：

# 加载CSV并按8:2比例拆分训练集和测试集
dataset = load_dataset('csv', data_files='/content/sequences.csv').train_test_split(test_size=0.2)

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)

# 此时可正常访问train和test拆分
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(1000))

方案2：手动拆分已加载的train数据集 #

如果已经加载了仅含train拆分的数据集，可手动拆分后重组：

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

# 将原train数据集拆分为新的train和test子集
split_dataset = dataset["train"].train_test_split(test_size=0.2)
# 重组为包含train和test的数据集字典
dataset = {
    "train": split_dataset["train"],
    "test": split_dataset["test"]
}

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))

额外提醒 #

• 确认CSV文件包含label列，否则后续训练会因找不到标签报错
• test_size参数可按需调整，比如设为0.1代表10%数据用作测试集

内容的提问来源于stack exchange，提问作者eneko valero