sdg/preprocessing/dedupe/test_semantic.py

"""Tests for SemanticDeduplicator.

Avoids downloading any embedding model by exercising `dedup_from_embeddings`
directly with handcrafted L2-normalized vectors. Skipped if faiss is not
installed.
"""

from __future__ import annotations

import pytest

pytest.importorskip("faiss")
import numpy as np

from sdg.preprocessing.dedupe.semantic import SemanticDeduplicator


def _normalize(v: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(v, axis=1, keepdims=True)
    return (v / np.clip(norms, 1e-12, None)).astype(np.float32)


def _vec(*components: float) -> np.ndarray:
    return np.array(components, dtype=np.float32)


# ── Boundary cases ───────────────────────────────────────────────────────────

def test_empty_returns_empty():
    dedup = SemanticDeduplicator()
    assert dedup.dedup_from_embeddings(np.zeros((0, 4), dtype=np.float32)) == []


def test_single_returns_single():
    dedup = SemanticDeduplicator()
    emb = _normalize(np.array([[1.0, 0.0, 0.0, 0.0]]))
    assert dedup.dedup_from_embeddings(emb) == [0]


# ── Identity / orthogonality ─────────────────────────────────────────────────

def test_identical_embeddings_collapse():
    dedup = SemanticDeduplicator(threshold=0.9)
    emb = _normalize(np.stack([_vec(1, 0, 0), _vec(1, 0, 0)]))
    keep = dedup.dedup_from_embeddings(emb)
    assert len(keep) == 1


def test_orthogonal_embeddings_both_kept():
    dedup = SemanticDeduplicator(threshold=0.5)
    emb = _normalize(np.stack([_vec(1, 0, 0), _vec(0, 1, 0)]))
    keep = dedup.dedup_from_embeddings(emb)
    assert keep == [0, 1]


# ── Threshold behavior ───────────────────────────────────────────────────────

def test_above_threshold_collapses():
    # cos(a, b) ≈ 0.95 (small angle)
    a = _vec(1.0, 0.0)
    b = _vec(np.cos(np.radians(18)), np.sin(np.radians(18)))  # ~0.951
    emb = _normalize(np.stack([a, b]))
    dedup = SemanticDeduplicator(threshold=0.92)
    keep = dedup.dedup_from_embeddings(emb)
    assert len(keep) == 1


def test_below_threshold_both_kept():
    # cos(a, b) ≈ 0.866 (30 degrees)
    a = _vec(1.0, 0.0)
    b = _vec(np.cos(np.radians(30)), np.sin(np.radians(30)))
    emb = _normalize(np.stack([a, b]))
    dedup = SemanticDeduplicator(threshold=0.92)
    keep = dedup.dedup_from_embeddings(emb)
    assert keep == [0, 1]


# ── Cluster of three ─────────────────────────────────────────────────────────

def test_cluster_of_three_via_chain_above_threshold():
    a = _vec(1.0, 0.0)
    b = _vec(np.cos(np.radians(10)), np.sin(np.radians(10)))
    c = _vec(np.cos(np.radians(20)), np.sin(np.radians(20)))
    emb = _normalize(np.stack([a, b, c]))
    dedup = SemanticDeduplicator(threshold=0.95)
    keep = dedup.dedup_from_embeddings(emb)
    # All three angles within ~20 degrees -> all pairwise cos > 0.93,
    # so they form one cluster -> one kept.
    assert len(keep) == 1


# ── Representative selection ─────────────────────────────────────────────────

def test_key_fn_selects_representative():
    emb = _normalize(np.stack([_vec(1, 0), _vec(1, 0), _vec(1, 0)]))
    response_lengths = [10, 500, 50]
    key_fn = lambda i: -response_lengths[i]
    dedup = SemanticDeduplicator(threshold=0.99)
    keep = dedup.dedup_from_embeddings(emb, key_fn=key_fn)
    assert keep == [1]


# ── HNSW path ────────────────────────────────────────────────────────────────

def test_hnsw_path_runs_above_threshold_size():
    """Force the HNSW code path with a tiny hnsw_threshold.

    Uses 384-dim random gaussians (matching our default embed model). At that
    dimensionality, random pairs have expected cosine ~0 with std ~1/sqrt(384),
    so a threshold of 0.95 is deep in the tail and nothing should collapse.
    """
    rng = np.random.default_rng(42)
    raw = rng.standard_normal((200, 384)).astype(np.float32)
    emb = _normalize(raw)
    dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=10)
    keep = dedup.dedup_from_embeddings(emb)
    assert len(keep) == 200


def test_hnsw_path_collapses_planted_duplicates():
    """HNSW should still find planted near-duplicates among many distractors."""
    rng = np.random.default_rng(123)
    distractors = rng.standard_normal((180, 384)).astype(np.float32)
    # Plant 20 copies of one vector
    seed_vec = rng.standard_normal((1, 384)).astype(np.float32)
    planted = np.repeat(seed_vec, 20, axis=0)
    raw = np.concatenate([distractors, planted], axis=0)
    emb = _normalize(raw)
    dedup = SemanticDeduplicator(threshold=0.95, hnsw_threshold=10, topk=25)
    keep = dedup.dedup_from_embeddings(emb)
    # 180 distractors + 1 representative of the planted cluster = 181
    assert len(keep) == 181


# ── Validation ───────────────────────────────────────────────────────────────

def test_invalid_threshold_raises():
    with pytest.raises(ValueError):
        SemanticDeduplicator(threshold=0.0)


def test_invalid_topk_raises():
    with pytest.raises(ValueError):
        SemanticDeduplicator(topk=0)


# ── Device resolution ────────────────────────────────────────────────────────

def test_device_explicit_passthrough():
    assert SemanticDeduplicator(device="cpu").resolve_device() == "cpu"


def test_device_auto_returns_known_string():
    """auto must resolve to one of mps/cuda/cpu (depending on host)."""
    resolved = SemanticDeduplicator(device="auto").resolve_device()
    assert resolved in {"mps", "cuda", "cpu"}