代码走读

这一页按真实文件顺序读项目。

src/preprocess.py

它负责把原始数据变成模型能吃的格式。

原始 TSV 大概是：

label    text_a
家居     樱之歌蓝色之恋5件套日式釉下彩...
家电     北欧欧慕料理机手持料理棒...

预处理做四件事：

dataset_dict = load_dataset(..., delimiter="\t")
all_labels = sorted(set(dataset_dict["train"]["label"]))
dataset_dict = dataset_dict.cast_column("label", ClassLabel(names=all_labels))
dataset_dict = dataset_dict.map(tokenize_batch, batched=True)
dataset_dict.save_to_disk(...)

最关键的是：

def tokenize_batch(examples):
    inputs = tokenizer(examples["text_a"], truncation=True)
    inputs["labels"] = examples["label"]
    return inputs

这里发生了两个转换：

text_a -> input_ids / attention_mask
label  -> labels

为什么字段叫 labels？

因为 HuggingFace 的分类模型看到 labels 后，会自动计算：

loss = CrossEntropyLoss(logits, labels)

src/train.py

train.py 是组装层。

它负责：

tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(...)
train_dataset = load_dataset_split("train")
valid_dataset = load_dataset_split("valid")
collate_fn = DataCollatorWithPadding(...)
trainer = Trainer(...)
trainer.train()

注意：它不是训练循环本身。

你可以把它理解成：

创建模型
创建数据
创建配置
创建 Trainer
启动训练

src/trainer.py

这个文件是训练核心。

Trainer.train()

它管全局节奏：

for epoch in range(self.config.epochs):
    self.model.train()

    for batch in train_dataloader:
        loss = self._train_one_step(batch)

        if self.global_step % self.config.save_steps == 0:
            metrics = self.evaluate()
            self._check_early_stop(metrics)
            self._save_checkpoint()

        self.global_step += 1

这里你要分清：

Trainer.train()      = 调度训练过程
_train_one_step()    = 训练一个 batch
evaluate()           = 跑验证集
_save_checkpoint()   = 保存断点续训状态
_save_best_model()   = 保存验证集表现最好的模型

_train_one_step()

项目里因为用了 AMP，代码是：

with torch.autocast(...):
    outputs = self.model(**batch)
    loss = outputs.loss

self.scaler.scale(loss).backward()
self.scaler.step(self.optimizer)
self.scaler.update()
self.optimizer.zero_grad()

你可以先脑内翻译成普通版：

outputs = model(**batch)
loss = outputs.loss

loss.backward()
optimizer.step()
optimizer.zero_grad()

AMP 是为了加速和省显存，不改变训练主逻辑。

evaluate()

验证函数的特点是：

self.model.eval()

with torch.no_grad():
    for batch in eval_dataloader:
        outputs = self.model(**batch)
        preds = torch.argmax(outputs.logits, dim=-1)

self.model.train()

它做的是：

关闭训练行为
不记录计算图
不反向传播
不更新参数
只统计 loss / acc / f1

最后重新调用：

self.model.train()

是因为 evaluate() 经常在训练过程中插入执行，评估完还要继续训练。

src/dataset.py

它负责把预处理后的数据从磁盘读出来：

dataset = load_from_disk(PROCESSED_DATA_DIR / split)

然后通过 DataLoader 按 batch 输出：

DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=shuffle,
    collate_fn=collate_fn,
)

collate_fn 这里通常是：

DataCollatorWithPadding(...)

它的作用是把一个 batch 内不同长度的文本动态补齐到一样长。

src/predict.py

预测阶段和训练阶段不同。

训练 batch 有：

input_ids
attention_mask
labels

预测时通常没有 labels：

input_ids
attention_mask

所以预测时模型一般只输出：

logits

然后：

pred_id = logits.argmax(dim=-1)
category = id2label[pred_id]

这就是把模型分数转成最终中文类别。