prerank_pipeline.py

1"""2粗排模块完整链路骨架（第 5 章 §5.1 为什么需要粗排 / 图 5.1 粗排模块内部处理流程）34对应正文概念5------------6正文 §5.1 指出粗排位于召回与精排之间，职责是「在保证较高召回率的前提下，利用轻量级7模型对候选集进行进一步筛选，将候选规模从成千上万压缩到数百的量级」，同时还要承担8多目标融合、流量调控（Quota / 保量）、候选集优化等任务。图 5.1「粗排模块内部处理9流程」给出的五段式链路是本文件的实现对象：1011    多路召回候选集  ->  粗排模型打分(多 pxtr)  ->  多目标融合  ->12    Quota 调控 + 保量  ->  截断 TopK  ->  输出给精排1314本实现做了什么15--------------161. 用一个 `PreRankPipeline` 把五个阶段串成可运行的 Python 链路，每一阶段都记录17   「进入候选量级 / 退出候选量级 / 耗时(ms) / 占用预算比例」，最后打印链路日志表，18   复现正文「粗排阶段能够获得的推理预算通常最多只有 50 毫秒」（§5.2）的预算视角。192. 打分、融合、Quota、保量、截断全部按正文公式实现：20   - 融合（§5.4.1 线性融合）：s_agg = w1*p_ctr + w2*p_cvr + w3*p_ltr21   - 保量（§5.6.2）：Exposure_i < Target_i  =>  Score_i = Score_i + Δ22   - Quota（§5.6.1）：CandidateSet = ∪_i TopK_i，配额按正文表 5.3 的 40/30/20/1023   - 截断（§5.6.3）：TopK(CandidateSet)，K=300243. 统计正文 §5.7 关心的候选集质量指标：类目覆盖数、最大类目占比（Candidate25   Concentration）、新内容（探索）占比。2627[MOCK] 组件28-----------29`MockRecallService`（多路召回）、`MockPreRankModel`（三塔 pxtr 打分）、30`MockRankService`（精排）、`MockExposureStore`（保量用的曝光计数存储）均为教学用假31实现，不访问网络、不读数据集，用确定性伪随机数生成分数。3233关键公式34--------35    s_agg = Σ_i w_i p_i                     （线性融合，§5.4.1）36    Score_i <- Score_i + Δ  if Exposure_i < Target_i   （保量，§5.6.2）37    CandidateSet = ∪_i TopK_i               （Quota，§5.6.1）38    FinalSet = TopK(CandidateSet), K = 300  （截断，§5.6.3）39"""4041import random42import time43from dataclasses import dataclass, field444546# ----------------------------------------------------------------------------47# 数据结构48# ----------------------------------------------------------------------------49@dataclass50class Item:51    """候选物品。pxtr 由粗排模型填充，score 由多目标融合填充。"""5253    item_id: int54    category: str                 # game / entertainment / ecommerce55    is_new_anchor: bool           # 是否新主播（新内容），用于新主播池与保量56    recall_source: str            # 来自哪一路召回57    pctr: float = 0.058    pcvr: float = 0.059    pltr: float = 0.060    score: float = 0.061    guarantee_bonus: float = 0.0  # 保量补偿项 Δ62    quota_pool: str = ""          # 最终被哪个 Quota 池选中636465@dataclass66class StageLog:67    """链路单阶段日志：候选量级 + 耗时。"""6869    name: str70    n_in: int71    n_out: int72    cost_ms: float73    note: str = ""747576# ----------------------------------------------------------------------------77# [MOCK] 外部依赖78# ----------------------------------------------------------------------------79class MockRecallService:80    """[MOCK] 多路召回服务。按正文 §5.1「多路召回策略」造 5 路候选并做去重融合。"""8182    SOURCES = {83        "u2i_cf": 1600,84        "i2i_swing": 1400,85        "vector_ann": 1200,86        "hot_pool": 500,87        "cold_start": 400,88    }89    CATEGORIES = ["game", "entertainment", "ecommerce"]9091    def __init__(self, rng: random.Random):92        self.rng = rng9394    def retrieve(self) -> tuple[list[Item], int]:95        """返回 (去重后的候选集, 去重前的召回总条数)。"""96        raw_total = 097        pool: dict[int, Item] = {}98        for source, n in self.SOURCES.items():99            raw_total += n100            for _ in range(n):101                # 物品 id 空间故意小于召回总量，从而产生多路重叠102                iid = self.rng.randrange(0, 6000)103                if iid in pool:104                    continue105                cat = self.CATEGORIES[iid % 3]106                is_new = (source == "cold_start") or (iid % 17 == 0)107                pool[iid] = Item(iid, cat, is_new, source)108        return list(pool.values()), raw_total109110111class MockPreRankModel:112    """[MOCK] 粗排多任务模型（§5.2）：一次前向同时输出 pctr / pcvr / pltr 三个 pxtr。113114    真实实现是双塔 + 多任务塔（代码见 `ch11 粗排模型/code/`），这里用物品 id 的115    确定性哈希生成分数，使同一 seed 下结果可复现；并刻意做两件事：116    (a) 三个 pxtr 量纲不同（pctr ~ 0.07、pcvr ~ 0.017、pltr ~ 0.25 量级），以呈现117        正文 §5.4.1「线性融合对于 pxtr 的绝对数值高度敏感」；118    (b) 给 game 类目一个正向偏置（模拟正文 §5.6.1 的重度游戏用户 P(Game|User)=0.9），119        使纯 TopK(score) 的候选集高度集中于游戏类，复现正文表 5.2「无 Quota 约束下120        的候选集示例」中 Top5 全是游戏的现象。121    """122123    # 单条候选的打分开销（微秒）。双塔只需一次低维点积，故极小。124    US_PER_ITEM = 4.0125    # 用户的类目兴趣先验，直接作用在 pctr 上（倍数）126    CATEGORY_PRIOR = {"game": 1.8, "entertainment": 1.0, "ecommerce": 0.7}127128    def predict_batch(self, items: list[Item]) -> None:129        for it in items:130            h = (it.item_id * 2654435761) % 1000003131            u = ((h % 997) + 0.5) / 997.0132            v = ((h // 997 % 991) + 0.5) / 991.0133            w = ((h // 7919 % 983) + 0.5) / 983.0134            prior = self.CATEGORY_PRIOR[it.category]135            it.pctr = round(min(0.99, (0.02 + 0.10 * u) * prior), 6)  # 点击率136            it.pcvr = round(0.002 + 0.03 * v, 6)                      # 转化率137            it.pltr = round(0.05 + 0.40 * w, 6)                       # 点赞率138139140class MockExposureStore:141    """[MOCK] 保量用的曝光计数存储（线上通常是 Redis / 实时指标服务）。"""142143    def __init__(self, rng: random.Random, target: int = 500):144        self.rng = rng145        self.target = target146        self._exposure: dict[int, int] = {}147148    def exposure_of(self, item_id: int) -> int:149        if item_id not in self._exposure:150            self._exposure[item_id] = self.rng.randrange(0, 900)151        return self._exposure[item_id]152153    def under_delivered(self, item_id: int) -> bool:154        """Exposure_i < Target_i 即判定为「欠投」，需要保量补偿。"""155        return self.exposure_of(item_id) < self.target156157158class MockRankService:159    """[MOCK] 下游精排服务。只用来呈现「精排单条开销远高于粗排」的预算差异。"""160161    US_PER_ITEM = 180.0  # 精排单条开销 ≈ 粗排的 45 倍162163    def estimated_cost_ms(self, n: int) -> float:164        return n * self.US_PER_ITEM / 1000.0165166167# ----------------------------------------------------------------------------168# 粗排链路169# ----------------------------------------------------------------------------170class PreRankPipeline:171    """图 5.1 的五段式粗排链路。"""172173    # §5.4.1 线性融合权重（与正文代码清单 5.1 的伪代码一致）174    FUSION_WEIGHTS = {"pctr": 0.5, "pcvr": 0.3, "pltr": 0.2}175    # §5.6.1 Quota 配置：正文假设粗排最终输出 100 个候选时为 40/30/20/10176    QUOTA_RATIO = {177        "game": 0.40,178        "entertainment": 0.30,179        "ecommerce": 0.20,180        "new_anchor": 0.10,181    }182    GUARANTEE_DELTA = 0.010   # §5.6.2 保量补偿项 Δ183    BUDGET_MS = 50.0          # §5.2「粗排阶段能够获得的推理预算通常最多只有 50 毫秒」184185    def __init__(self, seed: int = 42, truncate_k: int = 300):186        rng = random.Random(seed)187        self.recall = MockRecallService(rng)188        self.model = MockPreRankModel()189        self.exposure = MockExposureStore(rng)190        self.rank = MockRankService()191        self.truncate_k = truncate_k192        self.logs: list[StageLog] = []193194    # -- 阶段 1：候选集进入 ---------------------------------------------------195    def stage_recall(self) -> list[Item]:196        t0 = time.perf_counter()197        items, raw_total = self.recall.retrieve()198        cost = (time.perf_counter() - t0) * 1000.0199        self.logs.append(200            StageLog("1 多路召回融合", raw_total, len(items), cost,201                     f"5 路召回, 去重率 {1 - len(items) / raw_total:.1%}")202        )203        return items204205    # -- 阶段 2：粗排模型打分 ------------------------------------------------206    def stage_score(self, items: list[Item]) -> list[Item]:207        t0 = time.perf_counter()208        self.model.predict_batch(items)209        cost = (time.perf_counter() - t0) * 1000.0210        model_cost = len(items) * self.model.US_PER_ITEM / 1000.0211        self.logs.append(212            StageLog("2 粗排模型打分", len(items), len(items), cost,213                     f"3 个 pxtr, 等效线上开销 {model_cost:.1f}ms")214        )215        return items216217    # -- 阶段 3：多目标融合 --------------------------------------------------218    def stage_fuse(self, items: list[Item]) -> list[Item]:219        t0 = time.perf_counter()220        w = self.FUSION_WEIGHTS221        for it in items:222            it.score = w["pctr"] * it.pctr + w["pcvr"] * it.pcvr + w["pltr"] * it.pltr223        cost = (time.perf_counter() - t0) * 1000.0224        self.logs.append(225            StageLog("3 多目标融合(线性)", len(items), len(items), cost,226                     "s=0.5pctr+0.3pcvr+0.2pltr")227        )228        return items229230    # -- 阶段 4：保量 + Quota ------------------------------------------------231    def stage_guarantee(self, items: list[Item]) -> list[Item]:232        t0 = time.perf_counter()233        boosted = 0234        for it in items:235            if it.is_new_anchor and self.exposure.under_delivered(it.item_id):236                it.guarantee_bonus = self.GUARANTEE_DELTA237                it.score += it.guarantee_bonus238                boosted += 1239        cost = (time.perf_counter() - t0) * 1000.0240        self.logs.append(241            StageLog("4 保量补偿", len(items), len(items), cost,242                     f"{boosted} 条欠投新内容 Score+Δ({self.GUARANTEE_DELTA})")243        )244        return items245246    def stage_quota(self, items: list[Item], out_size: int) -> list[Item]:247        t0 = time.perf_counter()248        pools: dict[str, list[Item]] = {k: [] for k in self.QUOTA_RATIO}249        for it in items:250            if it.is_new_anchor:251                pools["new_anchor"].append(it)252            else:253                pools[it.category].append(it)254        selected: dict[int, Item] = {}255        detail = []256        for pool_name, ratio in self.QUOTA_RATIO.items():257            quota = int(round(out_size * ratio))258            bucket = sorted(pools[pool_name], key=lambda x: x.score, reverse=True)259            taken = bucket[:quota]260            for it in taken:261                it.quota_pool = pool_name262                selected[it.item_id] = it263            detail.append(f"{pool_name}:{len(taken)}/{quota}(池内{len(bucket)})")264        out = list(selected.values())265        cost = (time.perf_counter() - t0) * 1000.0266        self.logs.append(267            StageLog("5 Quota 调控", len(items), len(out), cost, " ".join(detail))268        )269        return out270271    # -- 阶段 5：截断 + 输出 -------------------------------------------------272    def stage_truncate(self, items: list[Item]) -> list[Item]:273        t0 = time.perf_counter()274        out = sorted(items, key=lambda x: x.score, reverse=True)[: self.truncate_k]275        cost = (time.perf_counter() - t0) * 1000.0276        self.logs.append(277            StageLog("6 截断 TopK", len(items), len(out), cost,278                     f"K={self.truncate_k}, 精排预估开销 "279                     f"{self.rank.estimated_cost_ms(len(out)):.0f}ms")280        )281        return out282283    def run(self, out_size: int = 400) -> list[Item]:284        items = self.stage_recall()285        items = self.stage_score(items)286        items = self.stage_fuse(items)287        items = self.stage_guarantee(items)288        items = self.stage_quota(items, out_size=out_size)289        return self.stage_truncate(items)290291292# ----------------------------------------------------------------------------293# 候选集质量指标（§5.7）294# ----------------------------------------------------------------------------295def candidate_set_report(items: list[Item]) -> dict[str, float]:296    n = max(1, len(items))297    cat_cnt: dict[str, int] = {}298    for it in items:299        cat_cnt[it.category] = cat_cnt.get(it.category, 0) + 1300    max_share = max(cat_cnt.values()) / n301    explore_rate = sum(1 for it in items if it.is_new_anchor) / n302    src_cnt = len({it.recall_source for it in items})303    return {304        "size": len(items),305        "category_coverage": len(cat_cnt),306        "max_category_share": max_share,307        "explore_rate": explore_rate,308        "recall_source_coverage": src_cnt,309        "avg_score": sum(it.score for it in items) / n,310    }311312313def print_pipeline_logs(logs: list[StageLog], budget_ms: float) -> None:314    total = sum(g.cost_ms for g in logs)315    print(f"{'阶段':<18}{'入':>7}{'出':>7}{'耗时ms':>9}{'占预算':>8}  说明")316    print("-" * 110)317    for g in logs:318        print(f"{g.name:<18}{g.n_in:>7}{g.n_out:>7}{g.cost_ms:>9.3f}"319              f"{g.cost_ms / budget_ms:>7.1%}  {g.note}")320    print("-" * 110)321    print(f"{'合计':<18}{'':>7}{'':>7}{total:>9.3f}{total / budget_ms:>7.1%}"322          f"  预算 {budget_ms:.0f}ms -> {'通过' if total <= budget_ms else '超预算'}")323324325if __name__ == "__main__":326    random.seed(42)327328    print("=" * 110)329    print("粗排完整链路：召回 -> 打分 -> 融合 -> 保量 -> Quota -> 截断 -> 精排（§5.1 图 5.1）")330    print("=" * 110)331332    pipe = PreRankPipeline(seed=42, truncate_k=300)333    final = pipe.run(out_size=300)334    print_pipeline_logs(pipe.logs, PreRankPipeline.BUDGET_MS)335336    print("\n[候选集质量指标 §5.7]")337    rep = candidate_set_report(final)338    for k, v in rep.items():339        print(f"  {k:<24}{v:.4f}" if isinstance(v, float) else f"  {k:<24}{v}")340341    print("\n[输出给精排的 Top-10]")342    print(f"{'rank':>5}{'item':>7}{'category':>15}{'pool':>14}"343          f"{'pctr':>9}{'pcvr':>9}{'pltr':>9}{'score':>10}{'Δ':>8}")344    for i, it in enumerate(final[:10], 1):345        print(f"{i:>5}{it.item_id:>7}{it.category:>15}{it.quota_pool:>14}"346              f"{it.pctr:>9.4f}{it.pcvr:>9.4f}{it.pltr:>9.4f}{it.score:>10.5f}"347              f"{it.guarantee_bonus:>8.3f}")348349    # ---- 对照：不做 Quota / 保量，纯 TopK(score) ----350    print("\n" + "=" * 110)351    print("对照实验：纯 TopK(score) vs 带 Quota+保量 的候选集结构（正文表 5.2 的问题）")352    print("=" * 110)353    plain = PreRankPipeline(seed=42, truncate_k=300)354    raw = plain.stage_recall()355    raw = plain.stage_score(raw)356    raw = plain.stage_fuse(raw)357    plain_top = sorted(raw, key=lambda x: x.score, reverse=True)[:300]358    rep_plain = candidate_set_report(plain_top)359    print(f"{'指标':<26}{'纯TopK':>12}{'Quota+保量':>14}")360    for k in ("category_coverage", "max_category_share", "explore_rate", "avg_score"):361        a, b = rep_plain[k], rep[k]362        print(f"{k:<26}{a:>12.4f}{b:>14.4f}")363    print("\n结论：纯 TopK 下候选集平均分更高但新内容占比与类目结构失衡；"364          "Quota+保量以少量平均分为代价换回生态结构（§5.6）。")365366    # ---- 截断 K 对精排算力的影响 ----367    print("\n" + "=" * 110)368    print("截断条数 K 与精排算力预算（§5.6.3）")369    print("=" * 110)370    rank_svc = MockRankService()371    print(f"{'K':>6}{'精排开销ms':>14}{'相对K=300':>12}")372    for k in (100, 300, 500, 1000, 5000):373        c = rank_svc.estimated_cost_ms(k)374        print(f"{k:>6}{c:>14.1f}{c / rank_svc.estimated_cost_ms(300):>12.2f}x")375    print("\n结论：若不截断直接把 5000 条召回送精排，精排开销 900ms，"376          "远超单次请求数百毫秒的总预算 —— 这正是粗排存在的根本原因（§5.1）。")