复旦NLP-Beginner任务三

任务：FudanNLP/nlp-beginner 的任务三。

label2id = {
    "entailment": 0,
    "neutral": 1,
    "contradiction": 2
}

import json

train_path = "snli_1.0/snli_1.0_train.jsonl"

examples = []
with open(train_path, "r", encoding="utf-8") as f:
    for i, line in enumerate(f):
        data = json.loads(line)
        label = data["gold_label"]
        s1 = data["sentence1"]
        s2 = data["sentence2"]

        if label == '-':
            continue

        examples.append((s1, s2, label2id[label]))

        if len(examples) <= 3:
            print(f"样例{i+1}")
            print(f"Label: {label}")
            print(f"Premise: {s1}")
            print(f"Hypothesis: {s2}")
            print("-" * 50)

    print(f"总共读取数据{len(examples)}个")

print("观察前三个样例")
for ex in examples[:3]:
    print(ex)

def tokenize(sentence):
    return sentence.lower().split()

from collections import Counter

def build_vocab(examples, min_freq=2):
    counter = Counter() # 创建
    for s1, s2, _ in tqdm(examples, desc="正在分词"):
        counter.update(tokenize(s1))
        counter.update(tokenize(s2))

    vocab = {"<PAD>": 0, "<UNK>": 1}
    for word, freq in tqdm(counter.items(), desc="过滤低频词"):
        if freq >= min_freq:
            vocab[word] = len(vocab)
    return vocab

def encode(sentence, vocab, max_len=30):
    tokens = tokenize(sentence)
    ids = [vocab.get(tok, vocab["<UNK>"]) for tok in tokens]
    if len(ids) < max_len:
        ids += [vocab["<PAD>"]] * (max_len - len(ids))
    else:
        ids = ids[:max_len]
    return ids

vocab = build_vocab(examples, min_freq=2)
print(f"词表大小为{len(vocab)}\n测试分词:原句、分词tokenize、编码word_id")
# 测试一个句子
print("A person on a horse jumps over a broken down airplane.")
print(tokenize("A person on a horse jumps over a broken down airplane."))
print(encode("A person on a horse jumps over a broken down airplane.", vocab))

from torch.utils.data import Dataset, DataLoader

class SNLIDataset(Dataset):
    def __init__(self, examples, vocab, max_len=30):
        self.examples = examples
        self.vocab = vocab
        self.max_len = max_len

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        s1, s2, label = self.examples[idx]
        s1_ids = encode(s1, self.vocab, self.max_len)
        s2_ids = encode(s2, self.vocab, self.max_len)
        return torch.tensor(s1_ids), torch.tensor(s2_ids), torch.tensor(label)

batch_size = 64

train_examples = examples[5000:] # 545k 条数据
val_examples = examples[:5000] # 5k 条数据

train_dataset = SNLIDataset(train_examples, vocab)
val_dataset = SNLIDataset(val_examples, vocab)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

for s1_batch, s2_batch, label_batch in train_loader:
    print("Premise batch shape:", s1_batch.shape)
    print("Hypothesis batch shape:", s2_batch.shape)
    print("Label batch shape", label_batch.shape)
    break

# 词嵌入 embedding
import torch.nn as nn

embed_dim = 100
vocab_size = len(vocab)

embedding = nn.Embedding(
    num_embeddings=vocab_size,
    embedding_dim=embed_dim,
    padding_idx=0
)
# padding_idx=0 表示 <PAD> 不参与梯度更新

# BiLSTM 编码器
# 对句子序列进行上下文编码
hidden_dim = 128

encoder = nn.LSTM(
    input_size=embed_dim,
    hidden_size=hidden_dim,
    batch_first=True,
    bidirectional=True
)

import torch
import torch.nn as nn
import torch.nn.functional as F

class BiAttention(nn.Module):
    """ token-to-token 双向注意力，带 PAD mask """
    def forward(self, x, y, x_mask, y_mask):
        """
        x: [B, Lx, H], y: [B, Ly, H]
        x_mask: [B, Lx]  (True为有效token)
        y_mask: [B, Ly]
        """
        # 相似度矩阵 e = x @ y^T -> [B, Lx, Ly]
        e = torch.matmul(x, y.transpose(1, 2))

        # 把 PAD 位置的注意力权重屏蔽掉（加 -inf）
        # 对 y 归一化时，需要 mask y 的 PAD
        y_mask_float = (~y_mask).unsqueeze(1).float()  # [B,1,Ly]  True表示PAD
        e_y = e.masked_fill(y_mask_float.bool(), float('-inf'))
        alpha = F.softmax(e_y, dim=2)  # x attend y

        # 对 x 归一化时，需要 mask x 的 PAD
        x_mask_float = (~x_mask).unsqueeze(2).float()  # [B,Lx,1]
        e_x = e.masked_fill(x_mask_float.bool(), float('-inf'))
        beta = F.softmax(e_x, dim=1)  # y attend x（注意：沿 Lx 归一化）

        # 对齐表示
        x_align = torch.matmul(alpha, y)                 # [B, Lx, H]
        y_align = torch.matmul(beta.transpose(1, 2), x)  # [B, Ly, H]
        return x_align, y_align

class ESIMLite(nn.Module):
    def __init__(self, vocab_size, embed_dim=100, hidden_dim=128, num_classes=3, padding_idx=0, dropout=0.2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
        self.encoder = nn.LSTM(embed_dim, hidden_dim, batch_first=True, bidirectional=True)

        self.attn = BiAttention()

        # ESIM 的增强表示： [x, x_align, x - x_align, x * x_align]
        self.proj = nn.Linear(4 * (2*hidden_dim), 2*hidden_dim)  # 降维
        self.composer = nn.LSTM(2*hidden_dim, hidden_dim, batch_first=True, bidirectional=True)

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Sequential(
            nn.Linear(8*hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, num_classes)
        )

    def masked_mean(self, x, mask):
        # x: [B, L, H], mask: [B, L]  True=有效
        mask = mask.unsqueeze(-1).float()               # [B,L,1]
        summed = torch.sum(x * mask, dim=1)             # [B,H]
        count = torch.clamp(mask.sum(dim=1), min=1e-6)  # [B,1]
        return summed / count

    def masked_max(self, x, mask):
        # 把 PAD 位置置为极小值，避免选到
        mask = mask.unsqueeze(-1)  # [B,L,1]
        x_masked = x.masked_fill(~mask, float('-inf'))
        return torch.max(x_masked, dim=1).values  # [B,H]

    def forward(self, s1, s2):
        """
        s1: [B, L], s2: [B, L]
        """
        # mask：非0即有效（0是<PAD>）
        s1_mask = (s1 != 0)
        s2_mask = (s2 != 0)

        # 1) 编码
        s1_emb = self.embedding(s1)  # [B,L,E]
        s2_emb = self.embedding(s2)

        s1_out, _ = self.encoder(s1_emb)  # [B,L,2H]
        s2_out, _ = self.encoder(s2_emb)

        # 2) 双向注意力
        s1_align, s2_align = self.attn(s1_out, s2_out, s1_mask, s2_mask)

        # 3) 增强表示
        f_s1 = torch.cat([s1_out, s1_align, s1_out - s1_align, s1_out * s1_align], dim=-1)
        f_s2 = torch.cat([s2_out, s2_align, s2_out - s2_align, s2_out * s2_align], dim=-1)

        # 降维 + 非线性
        f_s1 = F.relu(self.proj(f_s1))
        f_s2 = F.relu(self.proj(f_s2))

        # 4) 组合编码（composition）
        v1, _ = self.composer(f_s1)  # [B,L,2H]
        v2, _ = self.composer(f_s2)

        # 5) 池化（mean + max）
        v1_mean = self.masked_mean(v1, s1_mask)  # [B,2H]
        v1_max  = self.masked_max(v1,  s1_mask)  # [B,2H]
        v2_mean = self.masked_mean(v2, s2_mask)
        v2_max  = self.masked_max(v2,  s2_mask)

        v = torch.cat([v1_mean, v1_max, v2_mean, v2_max], dim=-1)  # [B, 8H]
        v = self.dropout(v)
        logits = self.classifier(v)  # [B, num_classes]
        return logits

import torch
from torch.optim import Adam
from tqdm import tqdm

# —— 设备选择 ——
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"using {device}")

# —— 实例化模型 ——
vocab_size = len(vocab)
model = ESIMLite(vocab_size=vocab_size, embed_dim=100, hidden_dim=128, num_classes=3, padding_idx=0, dropout=0.2)
model = model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=1e-3)

# train_loader, val_loader

def evaluate(model, loader):
    model.eval()
    total, correct, total_loss = 0, 0, 0.0
    with torch.no_grad():
        for s1_batch, s2_batch, label_batch in loader:
            s1_batch = s1_batch.to(device)
            s2_batch = s2_batch.to(device)
            label_batch = label_batch.to(device)

            logits = model(s1_batch, s2_batch)
            loss = criterion(logits, label_batch)
            total_loss += loss.item() * s1_batch.size(0)

            preds = logits.argmax(dim=-1)
            correct += (preds == label_batch).sum().item()
            total += s1_batch.size(0)
    return total_loss / total, correct / total

# —— 训练 ——
EPOCHS = 3
for epoch in range(1, EPOCHS+1):
    model.train()
    pbar = tqdm(train_loader, desc=f"Epoch {epoch}")
    for s1_batch, s2_batch, label_batch in pbar:
        s1_batch = s1_batch.to(device)
        s2_batch = s2_batch.to(device)
        label_batch = label_batch.to(device)

        optimizer.zero_grad()
        logits = model(s1_batch, s2_batch)
        loss = criterion(logits, label_batch)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)  # 稳定训练
        optimizer.step()

        pbar.set_postfix(loss=f"{loss.item():.4f}")

    val_loss, val_acc = evaluate(model, val_loader)
    print(f"[Val] loss={val_loss:.4f}  acc={val_acc:.4f}")

样例1
Label: neutral
Premise: A person on a horse jumps over a broken down airplane.
Hypothesis: A person is training his horse for a competition.
--------------------------------------------------
样例2
Label: contradiction
Premise: A person on a horse jumps over a broken down airplane.
Hypothesis: A person is at a diner, ordering an omelette.
--------------------------------------------------
样例3
Label: entailment
Premise: A person on a horse jumps over a broken down airplane.
Hypothesis: A person is outdoors, on a horse.
--------------------------------------------------
总共读取数据549367个
观察前三个样例
('A person on a horse jumps over a broken down airplane.', 'A person is training his horse for a competition.', 1)
('A person on a horse jumps over a broken down airplane.', 'A person is at a diner, ordering an omelette.', 2)
('A person on a horse jumps over a broken down airplane.', 'A person is outdoors, on a horse.', 0)
正在分词: 100%|██████████| 549367/549367 [00:01<00:00, 381546.35it/s]
过滤低频词: 100%|██████████| 56218/56218 [00:00<00:00, 5622207.49it/s]
词表大小为39414
测试分词:原句、分词tokenize、编码word_id
A person on a horse jumps over a broken down airplane.
['a', 'person', 'on', 'a', 'horse', 'jumps', 'over', 'a', 'broken', 'down', 'airplane.']
[2, 3, 4, 2, 5, 6, 7, 2, 8, 9, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
Premise batch shape: torch.Size([64, 30])
Hypothesis batch shape: torch.Size([64, 30])
Label batch shape torch.Size([64])
Epoch 1: 100%|██████████| 8506/8506 [01:34<00:00, 89.97it/s, loss=0.5522]
[Val] loss=0.5836  acc=0.7552
Epoch 2: 100%|██████████| 8506/8506 [01:33<00:00, 90.82it/s, loss=0.7230]
[Val] loss=0.5448  acc=0.7748
Epoch 3: 100%|██████████| 8506/8506 [01:33<00:00, 90.58it/s, loss=0.3730]
[Val] loss=0.5338  acc=0.7812