-- coding: utf-8 --
"""
train_langid_12.py ------ 12 类语种识别分类器(学习实验脚本)
任务:一句话 → 它是什么语言(12 分类)
模型:Embedding → RNN → 掩码 MaxPool → LayerNorm → Dropout → Linear → Softmax
损失:CrossEntropyLoss 优化器:Adam
数据:data/langid/v12(12 语种:ar en es hi ja ko pt ru th tr vi zh)
运行:
KMP_DUPLICATE_LIB_OK=TRUE python train_langid_12.py
输出:
train_log_12.json ------ 全部训练过程与结果,供 HTML 可视化读取
"""
import os
import csv
import json
import time
import random
import unicodedata
from collections import Counter, OrderedDict
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
─── 超参数 ────────────────────────────────────────────────────────────
SEED = 42
MAX_LEN = 40
EMBED_DIM = 64
HIDDEN_DIM = 64
DROPOUT = 0.3
BATCH_SIZE = 64
LR = 1e-3
EPOCHS = 10
NUM_THREADS = 1
HERE = os.path.dirname(os.path.abspath(file))
上面拼出来的路径在 Windows 上不对,直接写死
DATA_DIR = r".\data\langid\v12"
12 个语种:
LABELS = "ar", "en", "es", "hi", "ja", "ko", "pt", "ru", "th", "tr", "vi", "zh"
语种中文名(供 HTML 显示)
LABEL_NAMES = {
"ar": "阿拉伯语", "en": "英语", "es": "西班牙语", "hi": "印地语",
"ja": "日语", "ko": "韩语", "pt": "葡萄牙语", "ru": "俄语",
"th": "泰语", "tr": "土耳其语", "vi": "越南语", "zh": "中文",
}
─── 1. 文本归一化(必须和建数据时一致)────────────────────────────────
def normalize(text):
"""
NFKC 归一化 + 转小写 + 只保留「字母(L)」与「组合符号(M)」。
保留 M:印地语、泰语的元音是组合标记,只留字母会把它们砍成裸辅音。
"""
t = unicodedata.normalize("NFKC", text).lower()
return "".join(c for c in t if unicodedata.category©0 in ("L", "M"))
─── 2. 字符表构建 ─────────────────────────────────────────────────────
def build_char_vocab(sentences, min_freq=1):
"""
字符表 = 训练集里所有出现过的字符,每个字符发一个整数编号。
第 0 号 ,第 1 号 ,其余按出现次数从高到低编号。
"""
freq = Counter()
for s in sentences:
freq.update(normalize(s))
chars = sorted(((c, n) for c, n in freq.items() if n >= min_freq),
key=lambda kv: (-kv1, kv0))
vocab = OrderedDict(("", 0), ("", 1))
for ch, _ in chars:
vocabch = len(vocab)
return vocab, freq
def encode(text, vocab, maxlen=MAX_LEN):
ids = vocab.get(ch, 1) for ch in normalize(text):maxlen
ids += 0 * (maxlen - len(ids))
return ids
─── 3. 读数据 ─────────────────────────────────────────────────────────
def load_split(data_dir, name):
path = os.path.join(data_dir, f"{name}.csv")
rows = \[\]
with open(path, encoding="utf-8") as f:
for r in csv.DictReader(f):
rows.append((r"text", r"label"))
return rows
class TextDataset(Dataset):
def init (self, rows, vocab, label2id):
self.X = encode(t, vocab) for t, _ in rows
self.y = label2id\[l for _, l in rows]
def __len__(self):
return len(self.y)
def __getitem__(self, i):
return torch.tensor(self.X[i], dtype=torch.long), self.y[i]
─── 4. 模型 ───────────────────────────────────────────────────────────
class LangIDClassifier(nn.Module):
"""
Embedding → RNN → 掩码 MaxPool → LayerNorm → Dropout → Linear
MaxPool 前的 mask 让补齐位不参与池化:
补齐位喂进去的是全 0 向量,RNN 照样会往前走,
masked_fill 成 -inf 后取最大值就干净了。
"""
def __init__(self, vocab_size, num_classes,
embed_dim=EMBED_DIM, hidden_dim=HIDDEN_DIM, dropout=DROPOUT):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.rnn = nn.RNN(embed_dim, hidden_dim, batch_first=True)
self.norm = nn.LayerNorm(hidden_dim)
self.dropout = nn.Dropout(dropout)
self.fc = nn.Linear(hidden_dim, num_classes)
def forward(self, x):
mask = (x != 0).unsqueeze(-1) # (B, L, 1)
h, _ = self.rnn(self.embedding(x)) # (B, L, H)
h = h.masked_fill(~mask, -1e9)
pooled = h.max(dim=1)[0] # (B, H)
pooled = self.dropout(self.norm(pooled))
return self.fc(pooled) # (B, C)
─── 5. 评测工具 ───────────────────────────────────────────────────────
def predict(model, loader):
model.eval()
preds, golds = \[\], \[\]
with torch.no_grad():
for X, y in loader:
preds += model(X).argmax(dim=1).tolist()
golds += y.tolist()
return preds, golds
def macro_f1(preds, golds, num_classes):
f1s = \[\]
for c in range(num_classes):
tp = sum(p == c and g == c for p, g in zip(preds, golds))
fp = sum(p == c and g != c for p, g in zip(preds, golds))
fn = sum(p != c and g == c for p, g in zip(preds, golds))
if tp + fp + fn == 0:
continue
prec = tp / (tp + fp) if tp + fp else 0.0
rec = tp / (tp + fn) if tp + fn else 0.0
f1s.append(2 * prec * rec / (prec + rec) if prec + rec else 0.0)
return sum(f1s) / len(f1s)
def confusion_matrix(preds, golds, n):
mat = \[0 * n for _ in range(n)]
for p, g in zip(preds, golds):
matgp += 1
return mat
def per_class_accuracy(preds, golds, num_classes):
result = \[\]
for c in range(num_classes):
idx = i for i, g in enumerate(golds) if g == c
hit = sum(predsi == c for i in idx)
result.append({"correct": hit, "total": len(idx),
"acc": hit / max(1, len(idx))})
return result
─── 6. 训练 ───────────────────────────────────────────────────────────
def main():
torch.set_num_threads(NUM_THREADS)
random.seed(SEED)
torch.manual_seed(SEED)
train_rows = load_split(DATA_DIR, "train")
dev_rows = load_split(DATA_DIR, "dev")
test_rows = load_split(DATA_DIR, "test")
label2id = {l: i for i, l in enumerate(LABELS)}
num_classes = len(LABELS)
# 从训练集构建字符表(v12 没有预生成 vocab 文件)
vocab, freq = build_char_vocab([t for t, _ in train_rows])
vocab_size = len(vocab)
print(f"数据目录:{DATA_DIR}")
print(f" 训练 {len(train_rows)} 条 / 验证 {len(dev_rows)} 条 / 测试 {len(test_rows)} 条")
print(f" 类别 {num_classes} 个:{' '.join(LABELS)}")
print(f" 字符表 {vocab_size} 个字符(从训练集现算)")
print(f" 随机猜的准确率基准:{1/num_classes:.1%}")
train_loader = DataLoader(TextDataset(train_rows, vocab, label2id),
batch_size=BATCH_SIZE, shuffle=True)
dev_loader = DataLoader(TextDataset(dev_rows, vocab, label2id), batch_size=128)
test_loader = DataLoader(TextDataset(test_rows, vocab, label2id), batch_size=128)
model = LangIDClassifier(vocab_size, num_classes)
n_params = sum(p.numel() for p in model.parameters())
print(f" 模型参数量:{n_params:,}\n")
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
# ── 记录训练过程 ─────────────────────────────────────────────────
step_losses = [] # 每一步的 loss
epoch_records = [] # 每个 epoch 的汇总
epoch_boundaries = [] # 每个 epoch 在 step_losses 中的起止下标
t0 = time.time()
for epoch in range(1, EPOCHS + 1):
model.train()
epoch_start = len(step_losses)
total = 0.0
for X, y in train_loader:
logits = model(X)
loss = criterion(logits, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
step_losses.append(loss.item())
total += loss.item()
epoch_end = len(step_losses)
epoch_boundaries.append([epoch_start, epoch_end])
preds, golds = predict(model, dev_loader)
dev_acc = sum(p == g for p, g in zip(preds, golds)) / len(golds)
dev_f1 = macro_f1(preds, golds, num_classes)
train_loss_avg = total / len(train_loader)
print(f"Epoch {epoch:2d}/{EPOCHS} loss={train_loss_avg:.4f} "
f"dev_acc={dev_acc:.4f} dev_f1={dev_f1:.4f}")
epoch_records.append({
"epoch": epoch,
"train_loss": train_loss_avg,
"dev_acc": dev_acc,
"dev_f1": dev_f1,
})
elapsed = time.time() - t0
print(f"\n训练耗时:{elapsed:.0f} 秒")
# ── 测试集评测 ─────────────────────────────────────────────────────
preds, golds = predict(model, test_loader)
test_acc = sum(p == g for p, g in zip(preds, golds)) / len(golds)
test_f1 = macro_f1(preds, golds, num_classes)
cm = confusion_matrix(preds, golds, num_classes)
pca = per_class_accuracy(preds, golds, num_classes)
print(f"\n测试集:accuracy = {test_acc:.4f} macro-F1 = {test_f1:.4f}")
# 混淆最严重的组合
pairs = [(cm[g][p], LABELS[g], LABELS[p])
for g in range(num_classes) for p in range(num_classes) if g != p]
pairs.sort(reverse=True)
print(" 最常混淆的组合(真实 → 预测):")
for cnt, g, p in pairs[:5]:
if cnt == 0:
break
print(f" {g} → {p}: {cnt} 条")
print("\n 各语种准确率:")
for c, lab in enumerate(LABELS):
print(f" {lab} ({LABEL_NAMES[lab]:4s}): "
f"{pca[c]['correct']}/{pca[c]['total']} = {pca[c]['acc']:.3f}")
# ── 推理示例:每类取一条 ──────────────────────────────────────────
demos = []
seen_labels = set()
for text, label in test_rows:
if label not in seen_labels:
seen_labels.add(label)
demos.append((text, label))
if len(seen_labels) == num_classes:
break
demos.sort(key=lambda x: label2id[x[1]])
inference_examples = []
model.eval()
with torch.no_grad():
for text, gold in demos:
logits = model(torch.tensor([encode(text, vocab)]))
prob = torch.softmax(logits, dim=1)[0]
pid = int(prob.argmax())
top3 = torch.topk(prob, 3)
inference_examples.append({
"text": text,
"gold": gold,
"gold_name": LABEL_NAMES[gold],
"pred": LABELS[pid],
"pred_name": LABELS[pid],
"correct": pid == label2id[gold],
"confidence": float(prob[pid]),
"top3": [{"label": LABELS[int(i)],
"name": LABEL_NAMES[LABELS[int(i)]],
"prob": float(v)}
for i, v in zip(top3.indices, top3.values)],
})
print("\n--- 推理示例 ---")
for ex in inference_examples:
mark = "✓" if ex["correct"] else "✗"
print(f" [{mark} {ex['pred']} {ex['confidence']:.2f}] "
f"(真实:{ex['gold']}) {ex['text'][:50]}")
# ── 保存模型 ──────────────────────────────────────────────────────
ckpt_path = os.path.join(HERE, "langid_rnn_12.pt")
torch.save({"state_dict": model.state_dict(),
"vocab": dict(vocab),
"labels": LABELS,
"maxlen": MAX_LEN,
"embed_dim": EMBED_DIM,
"hidden_dim": HIDDEN_DIM,
"dropout": DROPOUT}, ckpt_path)
print(f"\n模型已保存:{ckpt_path}")
# ── 保存训练日志 JSON(供 HTML 读取)─────────────────────────────
log = {
"meta": {
"task": "12 类语种识别",
"model": "Embedding → RNN → MaxPool → LayerNorm → Dropout → Linear",
"labels": LABELS,
"label_names": LABEL_NAMES,
"num_classes": num_classes,
"vocab_size": vocab_size,
"num_params": n_params,
"train_size": len(train_rows),
"dev_size": len(dev_rows),
"test_size": len(test_rows),
"epochs": EPOCHS,
"batch_size": BATCH_SIZE,
"lr": LR,
"embed_dim": EMBED_DIM,
"hidden_dim": HIDDEN_DIM,
"dropout": DROPOUT,
"max_len": MAX_LEN,
"seed": SEED,
"elapsed_s": round(elapsed, 1),
"random_baseline": 1 / num_classes,
},
"step_losses": step_losses,
"epoch_boundaries": epoch_boundaries,
"epoch_records": epoch_records,
"test": {
"accuracy": test_acc,
"macro_f1": test_f1,
"confusion_matrix": cm,
"per_class": [
{"label": LABELS[c],
"name": LABEL_NAMES[LABELS[c]],
**pca[c]}
for c in range(num_classes)
],
"top_confusions": [
{"gold": g, "pred": p, "count": cnt}
for cnt, g, p in pairs[:8] if cnt > 0
],
},
"inference_examples": inference_examples,
}
log_path = os.path.join(HERE, "train_log_12.json")
with open(log_path, "w", encoding="utf-8") as f:
json.dump(log, f, ensure_ascii=False, indent=2)
print(f"训练日志已保存:{log_path}")
if name == "main ":
main()