微调DNABERT2执行trainer.train()遇_TensorSliceDataset不可下标错误

核心问题 #

Trainer是Hugging Face为PyTorch生态设计的训练工具，不直接兼容TensorFlow的tf.data.Dataset格式，这就是执行trainer.train()时抛出TypeError: '_TensorSliceDataset' object is not subscriptable的根本原因。

同时你的代码存在一个潜在问题：使用的AutoModel是基础特征提取模型，没有内置分类头，无法直接完成二分类任务的训练与预测。

修复步骤 #

1. 替换模型为带分类头的版本 #

将AutoModel改为AutoModelForSequenceClassification，指定二分类所需的num_labels=2：

from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(
    "zhihan1996/DNABERT-2-117M",
    num_labels=2,
    trust_remote_code=True
).to('cuda')

2. 转换为Trainer兼容的数据集格式 #

使用Hugging Face的datasets库创建兼容的Dataset对象（替代TensorFlow数据集）：

from datasets import Dataset

# 训练集转换
train_data = Dataset.from_dict({
    'input_ids': train_encodings['input_ids'],
    'attention_mask': train_encodings['attention_mask'],
    'labels': train_labels
})

# 测试集转换
test_data = Dataset.from_dict({
    'input_ids': test_encodings['input_ids'],
    'attention_mask': test_encodings['attention_mask'],
    'labels': test_labels
})

3. 更新Trainer初始化 #

将train_dataset和eval_dataset替换为上述转换后的数据集：

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=test_data,
    compute_metrics=compute_metrics,
)

完整修正后的代码 #

import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import Dataset
from transformers import TrainingArguments, Trainer
import numpy as np
import evaluate

# 加载模型与分词器
tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained(
    "zhihan1996/DNABERT-2-117M",
    num_labels=2,
    trust_remote_code=True
).to('cuda')

# 读取数据
df = pd.read_csv('/content/dev.csv', nrows=10)
df1 = pd.read_csv('/content/test.csv', nrows=10)

train_sequences = df.iloc[:, 0].tolist()
test_sequences = df1.iloc[:, 0].tolist()
train_labels = df.iloc[:, 1].tolist()
test_labels = df1.iloc[:, 1].tolist()

# 分词处理
train_encodings = tokenizer(train_sequences, truncation=True, padding=True)
test_encodings = tokenizer(test_sequences, truncation=True, padding=True)

# 转换为Trainer兼容的数据集
train_data = Dataset.from_dict({
    'input_ids': train_encodings['input_ids'],
    'attention_mask': train_encodings['attention_mask'],
    'labels': train_labels
})
test_data = Dataset.from_dict({
    'input_ids': test_encodings['input_ids'],
    'attention_mask': test_encodings['attention_mask'],
    'labels': test_labels
})

# 训练参数与评估指标
training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")
metric = evaluate.load("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 = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_data,
    eval_dataset=test_data,
    compute_metrics=compute_metrics,
)

trainer.train()

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