sentence_bert.py

1"""2Sentence-BERT（SBERT）模型（第10章 §10.3.3 基于预训练语义模型的向量召回）34核心思想：5- 原生 BERT 做句子相似度要用 Cross-Encoder：把两句拼成 [CLS] A [SEP] B [SEP] 联合编码。6  精度高，但一个 Query 要和 N 条候选各做一次前向，大规模检索复杂度 O(N) 次 Transformer7  推理，两两匹配更是 O(N^2)，完全无法满足在线召回的延迟要求。8- SBERT 借鉴孪生网络（Siamese Network）改成 Bi-Encoder：两句分别过**共享参数**的同一个9  编码器，各自 Pooling 得到定长句向量，再用余弦相似度比较：10      sim(u, v) = u^T v / (||u|| * ||v||)11  于是候选文本可以离线编码建 ANN 索引，在线只需编码 Query 一次 → O(1) 次推理 + 一次检索。12- Pooling 策略：CLS / Mean / Max，论文实验表明 Mean Pooling 最稳定，本实现采用 Mean Pooling。13- 训练目标：回归目标 MSE(cos(u,v), gold) 与检索常用的 Triplet Loss14      L = max(d(a, p) - d(a, n) + margin, 0)15  本文件两者都实现。1617论文出处：Reimers & Gurevych, Sentence-BERT: Sentence Embeddings using Siamese18BERT-Networks, EMNLP 2019.（PyTorch autograd 实现）19"""2021import math22import random2324import torch25import torch.nn as nn26import torch.nn.functional as F27import torch.optim as optim2829CLS = "[CLS]"303132class SentenceBERT(nn.Module):33    """34    孪生（Siamese）句向量编码器：两句共享同一套参数。3536    句向量： e = W_proj * MeanPool( TransformerEncoder( TokenEmb + PosEmb ) )37    使用 PyTorch autograd 完成端到端训练。38    """3940    def __init__(self, vocab: list, d_model: int = 16, n_heads: int = 2,41                 d_ff: int = 32, max_len: int = 8, lr: float = 0.05):42        super().__init__()43        self.vocab = vocab44        self.d_model = d_model45        self.max_len = max_len46        self.lr = lr47        self.tok2id = {w: i for i, w in enumerate(vocab)}4849        self.tok_emb = nn.Embedding(len(vocab), d_model)50        self.pos_emb = nn.Embedding(max_len, d_model)5152        encoder_layer = nn.TransformerEncoderLayer(53            d_model=d_model, nhead=n_heads, dim_feedforward=d_ff,54            dropout=0.0, batch_first=True, norm_first=False55        )56        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)57        self.proj = nn.Linear(d_model, d_model)5859    def _ids(self, tokens: list) -> list:60        return [self.tok2id[w] for w in ([CLS] + tokens)[: self.max_len]]6162    def encode(self, tokens: list) -> torch.Tensor:63        """返回 Mean Pooling 后的句向量 (d_model,)。"""64        ids = self._ids(tokens)65        T = len(ids)66        ids_t = torch.tensor([ids], dtype=torch.long)                 # (1, T)67        pos_t = torch.arange(T, dtype=torch.long).unsqueeze(0)        # (1, T)68        x = self.tok_emb(ids_t) + self.pos_emb(pos_t)                 # (1, T, d)69        h = self.encoder(x)                                           # (1, T, d)70        pooled = h[0].mean(dim=0)                                     # (d,)71        return self.proj(pooled)                                      # (d,)7273    def embed(self, tokens: list) -> list:74        """对外接口：只拿句向量（离线批量编码候选文本时用）。"""75        with torch.no_grad():76            return self.encode(tokens).tolist()7778    def train_regression(self, pairs: list, epochs: int = 20) -> None:79        """回归目标：L = ( cos(u, v) - gold )^2，让句向量余弦逼近人工标注的相似度。"""80        optimizer = optim.Adam(self.parameters(), lr=self.lr)81        data = list(pairs)82        for ep in range(epochs):83            random.shuffle(data)84            tot = 0.085            self.train()86            for sa, sb, gold in data:87                u = self.encode(sa)88                v = self.encode(sb)89                cos = F.cosine_similarity(u.unsqueeze(0), v.unsqueeze(0))90                gold_t = torch.tensor([[gold]], dtype=torch.float)91                loss = F.mse_loss(cos, gold_t.squeeze(1))92                optimizer.zero_grad()93                loss.backward()94                optimizer.step()95                tot += loss.item()96            if ep % 4 == 0 or ep == epochs - 1:97                print(f"    epoch {ep:2d}  mse_loss={tot / len(data):.4f}")9899    def triplet_loss(self, a: list, p: list, n: list, margin: float = 0.2) -> float:100        """Triplet Loss：L = max( d(a,p) - d(a,n) + margin, 0 )，d 取余弦距离 1-cos。"""101        with torch.no_grad():102            ea = torch.tensor(self.embed(a))103            ep_ = torch.tensor(self.embed(p))104            en = torch.tensor(self.embed(n))105            d_pos = 1 - F.cosine_similarity(ea.unsqueeze(0), ep_.unsqueeze(0)).item()106            d_neg = 1 - F.cosine_similarity(ea.unsqueeze(0), en.unsqueeze(0)).item()107            return max(d_pos - d_neg + margin, 0.0)108109110# ── 快速验证 ──────────────────────────────────────────────────────────────────111if __name__ == "__main__":112    random.seed(42)113    torch.manual_seed(42)114    # 合成语料：6 个商品主题簇，同簇内句子语义相近（gold=1），跨簇不相关（gold=0）115    clusters = [116        [["蓝牙", "耳机", "降噪"], ["无线", "耳机", "续航"], ["耳机", "降噪", "续航"]],117        [["电竞", "显示器", "高刷"], ["游戏", "显示器", "面板"], ["高刷", "面板", "显示器"]],118        [["男士", "纯棉", "短袖"], ["男士", "透气", "T恤"], ["纯棉", "透气", "短袖"]],119        [["智能", "手表", "运动"], ["运动", "手表", "跑步"], ["智能", "跑步", "手表"]],120        [["笔记本", "轻薄", "商务"], ["轻薄", "便携", "办公"], ["笔记本", "商务", "办公"]],121        [["保温杯", "不锈钢", "大容量"], ["保温杯", "便携", "大容量"], ["不锈钢", "大容量", "保温杯"]],122    ]123    vocab = [CLS] + sorted({w for cl in clusters for s in cl for w in s})124    model = SentenceBERT(vocab, d_model=16, n_heads=2, d_ff=32, max_len=8, lr=0.05)125126    pairs = []127    for ci, cl in enumerate(clusters):128        for a in range(len(cl)):129            pairs.append((cl[a], cl[(a + 1) % len(cl)], 1.0))130            for _ in range(2):131                cj = random.choice([j for j in range(len(clusters)) if j != ci])132                pairs.append((cl[a], random.choice(clusters[cj]), 0.0))133134    probe_pos = (["蓝牙", "耳机", "续航"], ["无线", "耳机", "降噪"])135    probe_neg = (["蓝牙", "耳机", "续航"], ["保温杯", "不锈钢", "大容量"])136137    def cos_sim(a, b):138        av = torch.tensor(model.embed(a))139        bv = torch.tensor(model.embed(b))140        return F.cosine_similarity(av.unsqueeze(0), bv.unsqueeze(0)).item()141142    before_p = cos_sim(*probe_pos)143    before_n = cos_sim(*probe_neg)144145    print(f"孪生编码器：vocab={len(vocab)}  d_model={model.d_model}  "146          f"头数=2  Pooling=Mean  训练对={len(pairs)}")147    print(f"\n微调前：同义句相似度={before_p:.4f}   不相关句相似度={before_n:.4f}")148    print("\n回归目标（MSE）微调：")149    model.train_regression(pairs, epochs=24)150151    after_p = cos_sim(*probe_pos)152    after_n = cos_sim(*probe_neg)153    print(f"\n微调后：同义句相似度={after_p:.4f}（{after_p - before_p:+.4f}）   "154          f"不相关句相似度={after_n:.4f}（{after_n - before_n:+.4f}）")155156    tl = model.triplet_loss(157        ["蓝牙", "耳机", "续航"], ["无线", "耳机", "降噪"], ["保温杯", "不锈钢", "大容量"]158    )159    print(f"Triplet Loss（anchor=蓝牙耳机续航, pos=无线耳机降噪, neg=保温杯）= {tl:.4f}")160161    corpus = [["蓝牙", "无线", "耳机"], ["电竞", "游戏", "高刷"], ["男士", "纯棉", "T恤"],162              ["智能", "手表", "跑步"], ["笔记本", "轻薄", "办公"], ["保温杯", "不锈钢", "便携"]]163    index = [(c, torch.tensor(model.embed(c))) for c in corpus]164    query = ["耳机", "降噪", "无线"]165    q_vec = torch.tensor(model.embed(query))166    hits = sorted(167        [(F.cosine_similarity(q_vec.unsqueeze(0), v.unsqueeze(0)).item(), c) for c, v in index],168        reverse=True169    )170    print(f"\n【语义召回】Query = {' '.join(query)}  （离线编码 {len(index)} 条，在线仅 1 次前向）")171    for s, c in hits[:3]:172        print(f"  {s:.4f}  {' '.join(c)}")173    print(f"\n对比：Cross-Encoder 需要 {len(corpus)} 次 Transformer 前向才能给出同样的排序；")174    print("      Bi-Encoder 把编码与匹配解耦，这正是工业界 Embedding 召回的标准做法。")