social-data-attribution-w2/tests/unlearning/test_sampling.py

"""Tests for stratified unlearning sampling."""

from __future__ import annotations

import json
import random

import pytest
from datasets import Dataset

from unlearning.data.sampling import (
    estimate_token_count,
    filter_by_min_tokens,
    sample_forget,
    sample_retain_stratified,
    save_sampling_manifest,
)


def _make_dataset(n_per_topic=100, topics=None, short_frac=0.3):
    """Build a mock HF dataset with controllable topic/length distribution."""
    if topics is None:
        topics = ["science", "arts", "sports", "health"]

    rows = {"doc_id": [], "text": [], "weborganizer_topic": []}
    for i, topic in enumerate(topics):
        for j in range(n_per_topic):
            idx = i * n_per_topic + j
            if j < int(n_per_topic * short_frac):
                text = " ".join(["word"] * 50)
            else:
                text = " ".join(["word"] * 600)
            rows["doc_id"].append(f"doc_{idx}")
            rows["text"].append(text)
            rows["weborganizer_topic"].append(topic)
    return Dataset.from_dict(rows)


class TestEstimateTokenCount:
    def test_empty_string(self):
        assert estimate_token_count("") == 0.0

    def test_known_count(self):
        text = "one two three four five"
        assert estimate_token_count(text) == pytest.approx(5 * 1.35)


class TestFilterByMinTokens:
    def test_drops_short_docs(self):
        ds = _make_dataset(n_per_topic=100, short_frac=0.5)
        filtered, coverage = filter_by_min_tokens(ds, min_tokens=512)
        assert len(filtered) < len(ds)
        for topic, stats in coverage.items():
            assert stats["after"] <= stats["before"]

    def test_keeps_all_when_threshold_zero(self):
        ds = _make_dataset(n_per_topic=50)
        filtered, _ = filter_by_min_tokens(ds, min_tokens=0)
        assert len(filtered) == len(ds)

    def test_coverage_stats_complete(self):
        topics = ["a", "b"]
        ds = _make_dataset(n_per_topic=20, topics=topics, short_frac=0.5)
        _, coverage = filter_by_min_tokens(ds, min_tokens=100)
        assert set(coverage.keys()) == set(topics)
        for stats in coverage.values():
            assert "before" in stats
            assert "after" in stats


class TestSampleForget:
    def test_single_topic(self):
        ds = _make_dataset(n_per_topic=200)
        rng = random.Random(42)
        texts, doc_ids = sample_forget(ds, ["science"], max_docs=50, rng=rng)
        assert len(texts) == 50
        assert len(doc_ids) == 50
        assert all(isinstance(t, str) for t in texts)

    def test_multi_topic_stratified(self):
        ds = _make_dataset(n_per_topic=200)
        rng = random.Random(42)
        texts, doc_ids = sample_forget(
            ds,
            ["science", "arts"],
            max_docs=100,
            rng=rng,
        )
        assert len(texts) == 100
        assert len(doc_ids) == 100

    def test_handles_sparse_topic(self):
        ds = _make_dataset(n_per_topic=5, short_frac=0.0)
        rng = random.Random(42)
        texts, _ = sample_forget(ds, ["science"], max_docs=100, rng=rng)
        assert len(texts) == 5

    def test_deterministic_with_same_seed(self):
        ds = _make_dataset(n_per_topic=200)
        t1, _ = sample_forget(ds, ["science"], 50, random.Random(42))
        t2, _ = sample_forget(ds, ["science"], 50, random.Random(42))
        assert t1 == t2


class TestSampleRetainStratified:
    def test_excludes_target(self):
        topics = ["a", "b", "c", "d"]
        ds = _make_dataset(n_per_topic=100, topics=topics, short_frac=0.0)
        rng = random.Random(42)
        texts, _, counts = sample_retain_stratified(
            ds,
            exclude_topics={"a"},
            docs_per_bin=10,
            rng=rng,
            all_topics=topics,
        )
        assert "a" not in counts
        assert set(counts.keys()) == {"b", "c", "d"}
        assert all(c == 10 for c in counts.values())
        assert len(texts) == 30

    def test_handles_sparse_bin(self):
        topics = ["big", "small"]
        rows = {
            "doc_id": [],
            "text": [],
            "weborganizer_topic": [],
        }
        for i in range(100):
            rows["doc_id"].append(f"big_{i}")
            rows["text"].append("word " * 600)
            rows["weborganizer_topic"].append("big")
        for i in range(3):
            rows["doc_id"].append(f"small_{i}")
            rows["text"].append("word " * 600)
            rows["weborganizer_topic"].append("small")
        ds = Dataset.from_dict(rows)
        rng = random.Random(42)
        _, _, counts = sample_retain_stratified(
            ds,
            exclude_topics=set(),
            docs_per_bin=50,
            rng=rng,
            all_topics=topics,
        )
        assert counts["big"] == 50
        assert counts["small"] == 3

    def test_deterministic(self):
        topics = ["a", "b", "c"]
        ds = _make_dataset(n_per_topic=100, topics=topics, short_frac=0.0)
        t1, _, _ = sample_retain_stratified(
            ds,
            {"a"},
            10,
            random.Random(42),
            all_topics=topics,
        )
        t2, _, _ = sample_retain_stratified(
            ds,
            {"a"},
            10,
            random.Random(42),
            all_topics=topics,
        )
        assert t1 == t2


class TestSaveManifest:
    def test_writes_json(self, tmp_path):
        path = save_sampling_manifest(
            str(tmp_path),
            ["d1", "d2"],
            ["d3", "d4"],
            seed=42,
            target_topics=["science"],
            config_snapshot={"min_tokens": 512},
        )
        assert path.exists()
        data = json.loads(path.read_text())
        assert data["seed"] == 42
        assert data["forget_count"] == 2
        assert data["retain_count"] == 2
        assert data["target_topics"] == ["science"]
        assert data["config"]["min_tokens"] == 512