CH 07 / SOURCE EXAMPLE

rerank_pipeline.py

ch07 重排模块/code/ch02_rerank_architecture/rerank_pipeline.py

这个示例讲什么

│ └── rerank_pipeline.py 五阶段重排链路:候选接入→序列生成→评估→打散→后处理

对应正文第 7 章 §7.2「重排模块架构」中的四个环节: 重排前处理(Pre-processing)→ 序列生成(List Generation)→ 序列评估(List Evaluation) → 多样性打散(Diversification)。 本实现做了什么 ---------------- 用纯 Python 串出一条端到端的重排链路,各环节都是可替换的组件: 1. [MOCK] `ForwardIndexService` / `RedisFeatureStore`:模拟正排索引服务与 Redis 在线特征存储, 提供物品属性(垂类/作者/时长/内容类型)与 item embedding; 2. `PreProcessor`:特征补全 + 属性初始化 + 数据清洗(丢弃缺 embedding、缺垂类、时长异常的候选); 3. `ListGenerator`:高斯采样多目标融合

上面来自章节 README 或源码自身的说明,不是自动验证的运行结果,也不能代替 PDF 的书中论断。讲解时请结合实际代码、输入输出与原书页码。

源码结构与阅读入口

  1. build_catalog · 第 46 行[MOCK] 物品底表:垂类按头部效应抽样,作者从该垂类的小作者池中抽。
  2. ForwardIndexService · 第 70 行[MOCK] 正排索引服务:按 item_id 批量返回静态属性。
  3. RedisFeatureStore · 第 86 行[MOCK] Redis 在线特征存储:返回 item embedding 与近期曝光计数。
  4. mock_rank_output · 第 104 行[MOCK] 精排模块输出:item_id + 各目标 pxtr 预估值。热门垂类整体分更高。
  5. PreProcessor · 第 117 行特征补全 + 属性初始化 + 数据清洗。
  6. ListGenerator · 第 164 行高斯采样融合权重 + 贪心排序生成多条候选序列。
  7. ListEvaluator · 第 215 行listwise 评估:按「是否继续下滑」建模曝光概率后加权求和。
  8. Diversifier · 第 252 行MMR 贪心打散:lambda*rel - (1-lambda)*max_sim。
  9. PostProcessor · 第 282 行相邻去重(垂类/作者)+ 内容类型窗口配额 + 运营强插。
  10. fmt_list · 第 324 行
  11. avg_pairwise_cos · 第 332 行
  12. diversity_report · 第 339 行
  13. main · 第 349 行

完整源码 · 455 行

1"""对应正文第 7 章 §7.2「重排模块架构」中的四个环节:2重排前处理(Pre-processing)→ 序列生成(List Generation)→ 序列评估(List Evaluation)3→ 多样性打散(Diversification)。45本实现做了什么6----------------7用纯 Python 串出一条端到端的重排链路,各环节都是可替换的组件:81. [MOCK] `ForwardIndexService` / `RedisFeatureStore`:模拟正排索引服务与 Redis 在线特征存储,9   提供物品属性(垂类/作者/时长/内容类型)与 item embedding;102. `PreProcessor`:特征补全 + 属性初始化 + 数据清洗(丢弃缺 embedding、缺垂类、时长异常的候选);113. `ListGenerator`:高斯采样多目标融合权重,贪心生成多条候选序列(§7.3);124. `ListEvaluator`:按「是否继续下滑」口径做 listwise 序列打分(§7.4);135. `Diversifier`:MMR 贪心打散(§7.5);146. `PostProcessor`:相邻去重 + 窗口配额 + 强插(§7.6)。1516关键公式(§7.2 定义 7.1)17--------------------------18    max_pi F(pi) = sum_{t=1}^{T} Rel(i_t, u) + lambda * Div(pi) + mu * Constrain(pi)19    可行解空间 O(P(N, T)) = O(N! / (N-T)!)2021每个环节都打印输入/输出规模与耗时占比,便于对照正文的流程图。22"""2324import math25import random26import time2728CATEGORIES = ("影视", "美食", "萌宠", "知识", "游戏")29CONTENT_TYPES = ("短视频", "短剧", "直播", "广告")30TARGETS = ("p_evr", "p_lvr", "p_ltr", "p_cmr", "p_ftr")31FUSION_MEAN = (0.2, 0.3, 0.2, 0.2, 0.1)3233EMB_DIM = 834LIST_LEN = 10353637# --------------------------------------------------------------------------38# [MOCK] 外部依赖39# --------------------------------------------------------------------------40# 真实候选池是「头部效应」的:热门垂类占比高,同垂类物品 embedding 聚簇、41# 精排分也整体偏高。只有这样打散环节才有实际作用。42CATEGORY_WEIGHTS = (0.42, 0.26, 0.14, 0.10, 0.08)43CATEGORY_SCORE_BONUS = {"影视": 0.10, "美食": 0.07, "萌宠": 0.03, "知识": 0.0, "游戏": 0.0}444546def build_catalog(n_items):47    """[MOCK] 物品底表:垂类按头部效应抽样,作者从该垂类的小作者池中抽。"""48    centroids = {}49    for c in CATEGORIES:50        v = [random.gauss(0.0, 1.0) for _ in range(EMB_DIM)]51        norm = math.sqrt(sum(x * x for x in v))52        centroids[c] = [x / norm for x in v]5354    catalog = {}55    for i in range(n_items):56        cat = random.choices(CATEGORIES, weights=CATEGORY_WEIGHTS)[0]57        catalog["item_%d" % i] = {58            "cat": cat,59            # 每个垂类只有 4 个头部作者,容易出现同作者相邻60            "author_id": "author_%s_%d" % (cat, random.randint(0, 3)),61            "duration": round(random.uniform(5.0, 300.0), 1),62            "content_type": random.choices(CONTENT_TYPES, weights=(0.70, 0.12, 0.10, 0.08))[0],63            # 同垂类物品围绕同一个质心,cos 相似度高64            "emb": [0.85 * centroids[cat][d] + 0.35 * random.gauss(0.0, 1.0)65                    for d in range(EMB_DIM)],66        }67    return catalog686970class ForwardIndexService:71    """[MOCK] 正排索引服务:按 item_id 批量返回静态属性。"""7273    def __init__(self, catalog, missing_rate=0.06):74        self._table = {75            k: {f: v[f] for f in ("cat", "author_id", "duration", "content_type")}76            for k, v in catalog.items()77            if random.random() >= missing_rate   # 模拟索引缺失,前处理需清洗78        }79        self.qps = 08081    def mget(self, item_ids):82        self.qps += 183        return {i: self._table.get(i) for i in item_ids}848586class RedisFeatureStore:87    """[MOCK] Redis 在线特征存储:返回 item embedding 与近期曝光计数。"""8889    def __init__(self, catalog, missing_rate=0.04):90        self._emb = {k: v["emb"] for k, v in catalog.items() if random.random() >= missing_rate}91        self._exposure = {92            k: random.choices([0, 1, 2, 5], weights=(0.8, 0.1, 0.06, 0.04))[0] for k in catalog93        }94        self.qps = 09596    def mget(self, item_ids):97        self.qps += 198        return {99            i: {"emb": self._emb.get(i), "recent_exposure": self._exposure.get(i, 0)}100            for i in item_ids101        }102103104def mock_rank_output(catalog):105    """[MOCK] 精排模块输出:item_id + 各目标 pxtr 预估值。热门垂类整体分更高。"""106    out = []107    for item_id, meta in catalog.items():108        bonus = CATEGORY_SCORE_BONUS[meta["cat"]]109        pxtr = {t: round(min(0.999, random.random() * 0.74 + bonus), 6) for t in TARGETS}110        out.append({"item_id": item_id, "pxtr": pxtr})111    return out112113114# --------------------------------------------------------------------------115# 1. 重排前处理116# --------------------------------------------------------------------------117class PreProcessor:118    """特征补全 + 属性初始化 + 数据清洗。"""119120    def __init__(self, index_service, feature_store):121        self.index = index_service122        self.store = feature_store123        self.drop_stats = {}124125    def run(self, rank_items):126        ids = [it["item_id"] for it in rank_items]127        attrs = self.index.mget(ids)       # 一次批量 RPC,而不是逐条128        feats = self.store.mget(ids)129        self.drop_stats = {"no_attr": 0, "no_emb": 0, "bad_duration": 0}130131        cleaned = []132        for it in rank_items:133            a = attrs.get(it["item_id"])134            if a is None:135                self.drop_stats["no_attr"] += 1136                continue137            f = feats.get(it["item_id"]) or {}138            emb = f.get("emb")139            if emb is None:140                self.drop_stats["no_emb"] += 1141                continue142            if not (1.0 <= a["duration"] <= 600.0):143                self.drop_stats["bad_duration"] += 1144                continue145            norm = math.sqrt(sum(v * v for v in emb)) or 1.0146            cleaned.append(147                {148                    "item_id": it["item_id"],149                    "pxtr": it["pxtr"],150                    "cat": a["cat"],151                    "author_id": a["author_id"],152                    "duration": a["duration"],153                    "content_type": a["content_type"],154                    "emb": [v / norm for v in emb],   # 归一化,供相似度计算155                    "recent_exposure": f.get("recent_exposure", 0),156                }157            )158        return cleaned159160161# --------------------------------------------------------------------------162# 2. 序列生成(§7.3)163# --------------------------------------------------------------------------164class ListGenerator:165    """高斯采样融合权重 + 贪心排序生成多条候选序列。166167    每组权重同时产出两个东西:168      - seq :前 list_len 个物品,作为交给评估器的「候选序列」;169      - pool:前 pool_len 个物品,作为后续多样性打散的选择空间170              (打散要有挑选余地,否则只能在固定集合里换顺序)。171    """172173    def __init__(self, mean=FUSION_MEAN, std=0.1, n_samples=8,174                 list_len=LIST_LEN, pool_len=30):175        self.mean = mean176        self.std = std177        self.n_samples = n_samples178        self.list_len = list_len179        self.pool_len = pool_len180181    def fuse(self, item, w):182        return sum(wi * item["pxtr"][t] for wi, t in zip(w, TARGETS))183184    def run(self, items):185        seqs = []186        for k in range(self.n_samples):187            # 第 0 组用基准权重(线上主策略),其余围绕基准做高斯扰动188            w = (189                list(self.mean)190                if k == 0191                else [m + random.gauss(0.0, self.std) for m in self.mean]192            )193            ranked = sorted(items, key=lambda it: -self.fuse(it, w))194            seqs.append(195                {196                    "weights": w,197                    "seq": ranked[: self.list_len],198                    "pool": ranked[: self.pool_len],199                    "fuse": {it["item_id"]: self.fuse(it, w) for it in ranked[: self.pool_len]},200                }201            )202        # 去掉完全重复的序列203        uniq, seen = [], set()204        for s in seqs:205            key = tuple(it["item_id"] for it in s["seq"])206            if key not in seen:207                seen.add(key)208                uniq.append(s)209        return uniq210211212# --------------------------------------------------------------------------213# 3. 序列评估(§7.4)214# --------------------------------------------------------------------------215class ListEvaluator:216    """listwise 评估:按「是否继续下滑」建模曝光概率后加权求和。"""217218    def __init__(self, p0=1.0, fatigue=0.75):219        self.p0 = p0220        self.fatigue = fatigue221222    @staticmethod223    def item_value(item):224        """单点分 f(x_i):取预估观看时长(秒)。"""225        return item["duration"] * (0.25 + 0.55 * item["pxtr"]["p_lvr"])226227    def p_slide(self, item, prev):228        p = 0.50 + 0.45 * item["pxtr"]["p_evr"]229        if prev is not None and prev["cat"] == item["cat"]:230            p *= self.fatigue231        if prev is not None and prev["author_id"] == item["author_id"]:232            p *= self.fatigue233        return min(p, 0.99)234235    def score(self, seq):236        total, prob, prev = 0.0, self.p0, None237        for item in seq:238            total += prob * self.item_value(item)239            prob *= self.p_slide(item, prev)240            prev = item241        return total242243    def run(self, candidate_seqs):244        scored = [(self.score(s["seq"]), s) for s in candidate_seqs]245        scored.sort(key=lambda x: -x[0])246        return scored247248249# --------------------------------------------------------------------------250# 4. 多样性打散(§7.5,MMR)251# --------------------------------------------------------------------------252class Diversifier:253    """MMR 贪心打散:lambda*rel - (1-lambda)*max_sim。"""254255    def __init__(self, lam=0.7):256        self.lam = lam257258    @staticmethod259    def cos(a, b):260        return sum(x * y for x, y in zip(a["emb"], b["emb"]))261262    def run(self, pool, rel_scores, k):263        """从 pool 中按 MMR 贪心选出 k 个物品(同时决定集合与顺序)。"""264        remaining = list(pool)265        selected = [max(remaining, key=lambda it: rel_scores[it["item_id"]])]266        remaining.remove(selected[0])267        while remaining and len(selected) < k:268            best, best_v = None, -1e18269            for cand in remaining:270                max_sim = max(self.cos(cand, s) for s in selected)271                v = self.lam * rel_scores[cand["item_id"]] - (1.0 - self.lam) * max_sim272                if v > best_v:273                    best, best_v = cand, v274            selected.append(best)275            remaining.remove(best)276        return selected277278279# --------------------------------------------------------------------------280# 5. 重排后处理(§7.6)281# --------------------------------------------------------------------------282class PostProcessor:283    """相邻去重(垂类/作者)+ 内容类型窗口配额 + 运营强插。"""284285    def __init__(self, max_same_cat_in_window=3, window=5, ad_quota=1, force_slot=3):286        self.max_same_cat = max_same_cat_in_window287        self.window = window288        self.ad_quota = ad_quota289        self.force_slot = force_slot290        self.actions = []291292    def run(self, seq, forced_item):293        out = []294        pool = list(seq)295        self.actions = []296        while pool:297            picked = None298            for cand in pool:299                if out and (out[-1]["cat"] == cand["cat"] or out[-1]["author_id"] == cand["author_id"]):300                    continue301                win = out[-(self.window - 1):] if self.window > 1 else []302                if sum(1 for x in win if x["cat"] == cand["cat"]) >= self.max_same_cat:303                    continue304                if cand["content_type"] == "广告" and \305                        sum(1 for x in out if x["content_type"] == "广告") >= self.ad_quota:306                    continue307                picked = cand308                break309            if picked is None:                   # 约束无法满足,放宽为取队首(兜底)310                picked = pool[0]311                self.actions.append("relax@pos%d:%s" % (len(out) + 1, picked["item_id"]))312            pool.remove(picked)313            out.append(picked)314315        if forced_item is not None:316            slot = min(self.force_slot - 1, len(out))317            out.insert(slot, forced_item)318            self.actions.append("force_insert@pos%d:%s" % (slot + 1, forced_item["item_id"]))319            out = out[:LIST_LEN]320        return out321322323# --------------------------------------------------------------------------324def fmt_list(seq):325    lines = []326    for i, it in enumerate(seq, 1):327        lines.append("  %2d. %-9s %-5s %-11s %-9s dur=%6.1fs"328                     % (i, it["item_id"], it["cat"], it["author_id"], it["content_type"], it["duration"]))329    return "\n".join(lines)330331332def avg_pairwise_cos(seq):333    pairs = [(a, b) for i, a in enumerate(seq) for b in seq[i + 1:]]334    if not pairs:335        return 0.0336    return sum(sum(x * y for x, y in zip(a["emb"], b["emb"])) for a, b in pairs) / len(pairs)337338339def diversity_report(seq, name):340    cats = {it["cat"] for it in seq}341    authors = {it["author_id"] for it in seq}342    adj_cat = sum(1 for a, b in zip(seq, seq[1:]) if a["cat"] == b["cat"])343    adj_author = sum(1 for a, b in zip(seq, seq[1:]) if a["author_id"] == b["author_id"])344    print("  %-10s 垂类覆盖=%d/%d 作者覆盖=%2d/%d 相邻同垂类=%d 相邻同作者=%d 平均两两cos=%.4f"345          % (name, len(cats), len(CATEGORIES), len(authors), len(seq),346             adj_cat, adj_author, avg_pairwise_cos(seq)))347348349def main():350    random.seed(42)351352    n_candidates = 300353    print("=" * 96)354    print("§7.2 重排模块四段链路:前处理 -> 序列生成 -> 序列评估 -> 多样性打散 -> 后处理")355    print("=" * 96)356357    catalog = build_catalog(n_candidates)358    index = ForwardIndexService(catalog)359    store = RedisFeatureStore(catalog)360    rank_items = mock_rank_output(catalog)361    cat_hist = {}362    for m in catalog.values():363        cat_hist[m["cat"]] = cat_hist.get(m["cat"], 0) + 1364    print("[精排输出] 候选 %d 个(每个带 %d 个目标的 pxtr 预估)" % (len(rank_items), len(TARGETS)))365    print("           候选池垂类分布(头部效应):%s"366          % ", ".join("%s=%d" % (c, cat_hist.get(c, 0)) for c in CATEGORIES))367368    timings = {}369370    # ---- 前处理 ----371    t0 = time.perf_counter()372    pre = PreProcessor(index, store)373    items = pre.run(rank_items)374    timings["前处理"] = time.perf_counter() - t0375    print("\n[1/5 前处理] %d -> %d 个候选;RPC 次数:正排 %d 次 / Redis %d 次"376          % (len(rank_items), len(items), index.qps, store.qps))377    print("           清洗丢弃:%s" % pre.drop_stats)378379    # ---- 序列生成 ----380    t0 = time.perf_counter()381    gen = ListGenerator()382    cand_seqs = gen.run(items)383    timings["序列生成"] = time.perf_counter() - t0384    print("\n[2/5 序列生成] 采样 %d 组融合权重 -> %d 条互不相同的候选序列(每条长度 %d)"385          % (gen.n_samples, len(cand_seqs), LIST_LEN))386    for k, s in enumerate(cand_seqs[:3]):387        print("           序列%d 权重=[%s] 头部=%s"388              % (k, ", ".join("%.3f" % w for w in s["weights"]),389                 "/".join(it["item_id"] for it in s["seq"][:4])))390391    # ---- 序列评估 ----392    t0 = time.perf_counter()393    evaluator = ListEvaluator()394    scored = evaluator.run(cand_seqs)395    timings["序列评估"] = time.perf_counter() - t0396    print("\n[3/5 序列评估] listwise 打分(曝光概率 x 单点分 求和):")397    for k, (v, s) in enumerate(scored):398        tag = "  <== 最优" if k == 0 else ""399        print("           序列%d value=%8.4f(点加和=%8.4f)%s"400              % (k, v, sum(evaluator.item_value(it) for it in s["seq"]), tag))401    best = scored[0][1]402    best_seq = best["seq"]403404    # ---- 多样性打散 ----405    t0 = time.perf_counter()406    # 相关性分 rel 用该组融合分归一化到 [0, 1],和相似度量级可比407    fmax = max(best["fuse"].values())408    rel = {k: v / fmax for k, v in best["fuse"].items()}409    div = Diversifier(lam=0.7)410    diversified = div.run(best["pool"], rel, LIST_LEN)411    timings["多样性打散"] = time.perf_counter() - t0412    print("\n[4/5 多样性打散] MMR(lambda=%.1f),从 %d 个候选中选 %d 个:"413          % (div.lam, len(best["pool"]), LIST_LEN))414    diversity_report(best_seq, "打散前")415    diversity_report(diversified, "打散后")416417    # ---- 后处理 ----418    t0 = time.perf_counter()419    forced = {420        "item_id": "op_video_0",421        "pxtr": {t: 0.5 for t in TARGETS},422        "cat": "知识",423        "author_id": "author_official",424        "duration": 60.0,425        "content_type": "运营",426        "emb": [0.0] * EMB_DIM,427        "recent_exposure": 0,428    }429    post = PostProcessor()430    final = post.run(diversified, forced)431    timings["后处理"] = time.perf_counter() - t0432    print("\n[5/5 后处理] 规则命中:%s" % (post.actions or "无"))433    diversity_report(final, "最终曝光")434435    print("\n最终曝光列表(长度 %d):" % len(final))436    print(fmt_list(final))437438    print("\n各环节耗时占比:")439    total = sum(timings.values())440    for k, v in timings.items():441        print("  %-12s %8.3f ms  (%5.1f%%)" % (k, v * 1000.0, v / total * 100.0))442    print("  %-12s %8.3f ms" % ("合计", total * 1000.0))443444    final_value = evaluator.score(final)445    print("\n[目标函数 F(pi)] 最优生成序列 value=%.4f -> 打散后 value=%.4f -> 后处理后 value=%.4f"446          % (scored[0][0], evaluator.score(diversified), final_value))447    print("                 相关性被让渡 %.2f%%,换来垂类覆盖 %d->%d、相邻同垂类 %d->%d。"448          % ((scored[0][0] - final_value) / scored[0][0] * 100.0,449             len({it['cat'] for it in best_seq}), len({it['cat'] for it in final}),450             sum(1 for a, b in zip(best_seq, best_seq[1:]) if a["cat"] == b["cat"]),451             sum(1 for a, b in zip(final, final[1:]) if a["cat"] == b["cat"])))452453454if __name__ == "__main__":455    main()

如何用于学习或编写新代码

  1. 先看章节 README 的依赖和运行成本,定位这个例子的输入、输出及关键函数。
  2. 在原书 PDF 中找到对应概念的具体页;代码演示与书中论断分开标注。
  3. 只在依赖和耗时可控时运行单个示例,修改一个参数并记录真实输出与失败情形。
  4. 参考实现新代码时说明复用、改动、测试与尚未完成的生产依赖;教学示例并非生产系统。