pipu-music/services/recommender/tests/test_train.py

"""ALS, неявные оценки, офлайн-метрики и правило активации модели."""

import numpy as np

from pipu_rec.train.als import train_als
from pipu_rec.train.metrics import evaluate, ndcg_at_k, recall_at_k
from pipu_rec.train.pipeline import holdout, should_activate
from pipu_rec.train.signals import DISLIKE, Interaction, event_weight, implicit_scores
from pipu_rec.wave.engine import CatalogIndex, Profile, Session, build_batch, model_version

DAY = 86400.0


def test_event_weights_follow_spec() -> None:
    assert event_weight("play_end", 200_000, 200_000, True) == 1.0
    assert event_weight("skip", 120_000, 200_000, True) == 0.7  # ≥ 50% — засчитано
    assert event_weight("skip", 5_000, 200_000, True) == -1.0
    assert event_weight("skip", 20_000, 200_000, True) == -0.7
    assert event_weight("skip", 5_000, 200_000, False) == -0.5  # вне волны — слабее
    assert event_weight("play_start", 0, 200_000, True) is None


def test_implicit_scores_decay_and_dislike() -> None:
    now = 1000 * DAY
    inter = [
        Interaction("u", "a", now, 1.0),
        Interaction("u", "b", now - 90 * DAY, 1.0),  # полупериод — половина веса
        Interaction("u", "c", now, 3.0),
        Interaction("u", "c", now, DISLIKE),  # дизлайк перекрывает
    ]
    s = implicit_scores(inter, now)["u"]
    assert s["a"] == 1.0
    assert abs(s["b"] - 0.5) < 1e-9
    assert s["c"] == DISLIKE


def two_taste_world() -> dict[str, dict[str, float]]:
    """Две группы слушателей с непересекающимися вкусами: rock_* и jazz_*."""
    rng = np.random.default_rng(1)
    scores: dict[str, dict[str, float]] = {}
    for g in ("rock", "jazz"):
        for u in range(15):
            liked = rng.choice(30, size=12, replace=False)
            scores[f"{g}_user{u}"] = {f"{g}_{i}": 1.0 + float(rng.random()) for i in liked}
    return scores


def test_als_learns_taste_groups() -> None:
    model = train_als(two_taste_world(), factors=8, iterations=10)
    assert model is not None
    s = model.scores("rock_user0")
    rock = np.mean([s[i] for i, t in enumerate(model.items) if t.startswith("rock")])
    jazz = np.mean([s[i] for i, t in enumerate(model.items) if t.startswith("jazz")])
    assert rock > jazz + 0.3


def test_metrics() -> None:
    assert ndcg_at_k(["a", "b"], {"a"}) == 1.0
    assert 0 < ndcg_at_k(["x", "a"], {"a"}) < 1
    assert recall_at_k(["a", "x"], {"a", "b"}) == 0.5
    res = evaluate(["u", "v"], {"u": {"a"}}, lambda u: ["a"])
    assert res == {"users": 1, "ndcg@10": 1.0, "recall@50": 1.0}


def test_holdout_keeps_only_new_positives() -> None:
    now = 100 * DAY
    inter = [
        Interaction("u", "old", now - 30 * DAY, 1.0),
        Interaction("u", "old", now - DAY, 1.0),  # уже был в обучении — не «новый»
        Interaction("u", "new", now - DAY, 1.0),
        Interaction("u", "skipped", now - DAY, -1.0),
    ]
    train, relevant = holdout(inter, now)
    assert set(train["u"]) == {"old"}
    assert relevant == {"u": {"new"}}


def test_activation_gate() -> None:
    cur = {"offline": {"als": {"ndcg@10": 0.30}}}
    assert should_activate({"als": {"ndcg@10": 0.295}}, cur)  # −1.7% — допустимо
    assert not should_activate({"als": {"ndcg@10": 0.28}}, cur)
    assert should_activate({"users": 0}, cur)  # не на чем проверить
    assert should_activate({"als": {"ndcg@10": 0.1}}, None)


def test_engine_uses_cf_source() -> None:
    n = 40
    ids = [f"t{i}" for i in range(n)]
    cf = [
        np.array([1.0, 0.0], dtype=np.float32) if i < 20 else np.array([0.0, 1.0], dtype=np.float32)
        for i in range(n)
    ]
    index = CatalogIndex.build(
        ids=ids, artist=[f"a{i}" for i in range(n)], raw_emb=[None] * n, features={}, plays=[0] * n, raw_cf=cf
    )
    profile = Profile(
        scores={},
        disliked=set(),
        disliked_artists=set(),
        recent=set(),
        played=set(),
        liked_recent=set(),
        rediscover=set(),
        exploration=0.0,
        taste_cf=[1.0, 0.0],
    )
    picks = build_batch(index, profile, Session(), [], n=5, rng=np.random.default_rng(0))
    assert len(picks) == 5
    assert sum(int(p.id[1:]) < 20 for p in picks) >= 4  # почти всё — из «своей» половины по ALS
    assert model_version(index, profile).endswith("+cf")


def test_cf_vectors_are_padded_to_column_size() -> None:
    from pipu_rec.train.pipeline import CF_DIM, pad

    v = np.ones((3, 5), dtype=np.float32)
    assert pad(v).shape == (3, CF_DIM)
    assert pad(np.ones((2, 80), dtype=np.float32)).shape == (2, CF_DIM)


def test_bandit_params_and_exploration() -> None:
    from pipu_rec.train.bandit import compute, exploration

    now = 100 * DAY
    rows = [("u", "fresh", "complete", now)] * 8 + [("u", "fresh", "skip", now)] * 2
    rows += [("u", "audio_sim", "complete", now)] * 2 + [("u", "audio_sim", "skip_early", now)] * 8
    rows += [("u", "cf", "complete", now - 69 * DAY)]  # затухание: ~0.5
    p = compute(rows, now)["u"]
    assert p["fresh"] == [9.0, 3.0]
    assert p["audio_sim"] == [3.0, 9.0]
    assert 1.4 < p["cf"][0] < 1.6
    assert exploration(p) > 0.15  # «новое» заходит лучше остального — исследования больше
    assert exploration({"fresh": [1.0, 9.0], "audio_sim": [9.0, 1.0]}) == 0.05


def test_bandit_favours_successful_source() -> None:
    n = 60
    ids = [f"t{i}" for i in range(n)]
    index = CatalogIndex.build(
        ids=ids, artist=[f"a{i}" for i in range(n)], raw_emb=[None] * n, features={}, plays=[0] * n
    )
    good = Profile(
        scores={},
        disliked=set(),
        disliked_artists=set(),
        recent=set(),
        played=set(),
        liked_recent=set(),
        rediscover=set(),
        exploration=0.0,
        bandit={"popular": [50.0, 1.0], "explore": [1.0, 50.0]},
    )
    picks = build_batch(index, good, Session(), [], n=5, rng=np.random.default_rng(1))
    assert sum(p.source == "popular" for p in picks) >= 4