import unittest from peach import related from peach.related import rank_related def _closeness(left, right, weights): return related._closeness_features(related._feature(left), related._feature(right), weights) def _reasons(source, candidate): return related._reasons_features(related._sets(source), related._sets(candidate)) def _naive_rank(source, candidates, limit, *, seed=None, diversity=0.22): """优化前的朴素 MMR：每选一条，就把整池和**全部**已选重算一遍 closeness。 它留在测试里只有一个用途：当基准。增量维护「候选对已选集的最大 closeness」是等价变换，一旦哪天有人把它改成不等价的写法（比如改成只和 最后一条比、或者改成平均值），推荐顺序会静静地变，而不会报错。 """ pool = [dict(candidate) for candidate in candidates if int(candidate["id"]) != int(source["id"])] weights = related._idf([source, *pool]) for candidate in pool: candidate["_relevance"] = _closeness(source, candidate, weights) candidate["_reasons"] = _reasons(source, candidate) pool = [candidate for candidate in pool if candidate["_relevance"] > 0 and candidate["_reasons"]] selected = [] stable_seed = seed or str(source["id"]) while pool and len(selected) < max(0, int(limit)): def key(candidate): redundancy = max( (_closeness(candidate, prior, weights) for prior in selected), default=0.0, ) mmr = float(candidate["_relevance"]) - diversity * redundancy return (-mmr, related._stable_key(stable_seed, int(candidate["id"]))) choice = min(pool, key=key) choice["ranking_score"] = round(-key(choice)[0], 4) pool.remove(choice) reasons = choice.pop("_reasons") relevance = choice.pop("_relevance") choice["why"] = " · ".join(reasons[:2]) choice["shared_tags"] = sorted(related._sets(source)["tag"] & related._sets(choice)["tag"], key=lambda tag: (-weights.get(tag, 1.0), tag)) choice["related_score"] = round(relevance, 4) selected.append(choice) return selected def _corpus(size=120): """往每个维度都撑出重叠的候选池。模数互质，于是同创作者、同厂牌、 标签重叠和时长接近会交错出现，MMR 的冗余惩罚才真的会改变名次。""" rows = [] for asset_id in range(2, 2 + size): rows.append({ "id": asset_id, "entities": { "tag": {10 + asset_id % 7, 20 + asset_id % 5, 30 + asset_id % 11}, "creator": {40 + asset_id % 3} if asset_id % 4 else set(), "performer": {50 + asset_id % 6}, "studio": {60 + asset_id % 2}, "series": {70} if asset_id % 9 == 0 else set(), }, "year": 2010 + asset_id % 12, "duration": 600 + 37 * (asset_id % 20), }) return rows def _corpus_source(): return { "id": 1, "entities": {"tag": {10, 20, 30}, "creator": {40}, "performer": {50}, "studio": {60}, "series": {70}}, "year": 2016, "duration": 1200, } class RelatedRankerTests(unittest.TestCase): def test_actual_mmr_score_contains_redundancy_penalty(self): source = {"id": 1, "entities": {"tag": {10}}} candidates = [{"id": aid, "entities": {"tag": {10}}} for aid in (2, 3)] first, second = rank_related(source, candidates, 2) self.assertEqual(first["ranking_score"], first["related_score"]) self.assertAlmostEqual(second["ranking_score"], second["related_score"] - 0.22 * 0.55) def test_field_sort_covers_the_pool_and_keeps_missing_values_last(self): source = {"id": 1, "entities": {"tag": {10}}} candidates = [{"id": aid, "entities": {"tag": {10}}, "source_score": score} for aid, score in ((2, 1.0), (3, None), (4, 4.5))] for descending, expected in ((True, [4, 2, 3]), (False, [2, 4, 3])): rows = rank_related(source, candidates, 3, order_by="source_score", descending=descending) self.assertEqual([row["id"] for row in rows], expected) def test_rare_shared_tag_beats_common_shared_tag(self): source = {"id": 1, "entities": {"tag": {10, 20}}} common = {"id": 2, "entities": {"tag": {10}}} rare = {"id": 3, "entities": {"tag": {20}}} background = [ {"id": asset_id, "entities": {"tag": {10}}} for asset_id in range(4, 14) ] ranked = rank_related(source, [common, rare, *background], 2, seed="fixed") self.assertEqual(ranked[0]["id"], 3) self.assertEqual(ranked[0]["why"], "标签接近") def test_shared_tags_come_rarest_first(self): source = {"id": 1, "entities": {"tag": {10, 20, 30}}} both = {"id": 2, "entities": {"tag": {10, 20, 30}}} background = [{"id": asset_id, "entities": {"tag": {10, 30} if asset_id % 2 else {10}}} for asset_id in range(3, 13)] ranked = rank_related(source, [both, *background], 1, seed="fixed") self.assertEqual(ranked[0]["shared_tags"], [20, 30, 10]) def test_result_is_stable_and_mmr_adds_variety(self): source = {"id": 1, "entities": {"creator": {7}, "tag": {10, 11}}} candidates = [ {"id": 2, "entities": {"creator": {7}, "tag": {10, 11}}}, {"id": 3, "entities": {"creator": {7}, "tag": {10, 11}}}, {"id": 4, "entities": {"creator": {7}, "tag": {10}}}, ] first = rank_related(source, candidates, 3, seed="fixed", diversity=0.4) second = rank_related(source, candidates, 3, seed="fixed", diversity=0.4) self.assertEqual([row["id"] for row in first], [row["id"] for row in second]) self.assertEqual(first[0]["why"], "同创作者 · 标签接近") def test_unrelated_items_are_excluded(self): source = {"id": 1, "entities": {"tag": {10}}} ranked = rank_related(source, [{"id": 2, "entities": {"tag": {99}}}], 10) self.assertEqual(ranked, []) def test_incremental_mmr_picks_exactly_what_the_naive_loop_picks(self): """这次优化只能改代价，不能改结果。 逐条对整个返回字典，不只对 id：`why` 和 `related_score` 同样是契约， 它们直接显示在“接着看”和 Mix 队列里。diversity=0 把 MMR 退化成纯排序， 0.6 则让冗余惩罚大到足以改变名次；两端都要对得上才算等价。 """ source, corpus = _corpus_source(), _corpus() for limit in (0, 1, 3, 8, 25, 200): for diversity in (0.0, 0.22, 0.6): with self.subTest(limit=limit, diversity=diversity): expected = _naive_rank(source, corpus, limit, seed="fixed", diversity=diversity) actual = rank_related(source, corpus, limit, seed="fixed", diversity=diversity) self.assertEqual(actual, expected) def test_mmr_cost_is_linear_in_limit(self): """守住优化的实际目的：代价对 limit 是线性，不是二次。 旧写法在真实候选池（上限 4000）上 limit=28 要 6.2 秒、limit=60 要 28 秒， 而悬浮 Mix 卡片和详情页“接着看”都卡在这一步。上界取 池×(limit+2)： 增量实现宽松成立，任何「每轮重算全部已选」的写法当场超出。 """ source, corpus = _corpus_source(), _corpus() calls = 0 original = related._closeness_features def counting(left, right, weights): nonlocal calls calls += 1 return original(left, right, weights) related._closeness_features = counting try: picked = rank_related(source, corpus, 16, seed="fixed") finally: related._closeness_features = original self.assertEqual(len(picked), 16) self.assertLessEqual(calls, len(corpus) * 18, f"closeness 调用 {calls} 次，已超出线性上界") if __name__ == "__main__": unittest.main()