embedding.py

"""
embedding.py — Semantic similarity matching using sentence embeddings.

Compares a book's title+author against Audible search results using
cosine similarity. Runs on CUDA when available, falls back to CPU.

The model is loaded once and cached for the process lifetime.

.env:
    EMBEDDING_MODEL=all-MiniLM-L6-v2    # any sentence-transformers model
    EMBEDDING_THRESHOLD=0.75             # min cosine similarity to auto-match
    EMBEDDING_DEVICE=                    # "cuda", "cpu", or empty for auto
"""

import logging
import os
from typing import Optional

log = logging.getLogger(__name__)

_model = None
_device = None


def _load_model(model_name: str, device: str = ""):
    """Load sentence-transformers model once, cache globally."""
    global _model, _device
    if _model is not None:
        return

    try:
        from sentence_transformers import SentenceTransformer
        import torch
    except ImportError as e:
        log.warning(f"  sentence-transformers/torch import failed: {e}")
        _model = False  # Sentinel: tried and failed
        return
    except Exception as e:
        log.warning(f"  Unexpected error importing embedding deps: {e}")
        _model = False
        return

    if not device:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    _device = device

    log.info(f"  Loading embedding model '{model_name}' on {device}")

    # Suppress noisy HuggingFace HTTP logs and torch weight warnings
    import warnings
    warnings.filterwarnings("ignore", message=".*position_ids.*")
    logging.getLogger("sentence_transformers").setLevel(logging.WARNING)
    logging.getLogger("httpx").setLevel(logging.WARNING)
    logging.getLogger("huggingface_hub").setLevel(logging.WARNING)

    _model = SentenceTransformer(model_name, device=device)


def _build_text(title: str, author: str) -> str:
    """Build a comparison string from title + author."""
    parts = [title]
    if author:
        parts.append(f"by {author}")
    return " ".join(parts)


def _result_text(r: dict) -> str:
    """Build a comparison string from an Audible result."""
    parts = [r.get("title", "")]
    if r.get("author"):
        parts.append(f"by {r['author']}")
    if r.get("narrator"):
        parts.append(f"narrated by {r['narrator']}")
    if r.get("series"):
        parts.append(f"({r['series']})")
    return " ".join(parts)


def pick_best_match(
    title: str,
    author: str,
    results: list,
    model_name: str = "all-MiniLM-L6-v2",
    threshold: float = 0.75,
    device: str = "",
) -> tuple[Optional[int], float]:
    """
    Score results by cosine similarity against title+author.

    Returns (best_index, best_score).
    best_index is None if best score < threshold or model unavailable.
    """
    if not results:
        return None, 0.0

    _load_model(model_name, device)

    if _model is None or _model is False:
        return None, 0.0

    query = _build_text(title, author)
    candidates = [_result_text(r) for r in results]

    # Encode everything in one batch
    all_texts = [query] + candidates
    embeddings = _model.encode(
        all_texts, convert_to_tensor=True, show_progress_bar=False)

    # Cosine similarity between query and each candidate
    from sentence_transformers.util import cos_sim
    scores = cos_sim(embeddings[0:1], embeddings[1:])[0]

    # Find best
    best_idx = int(scores.argmax())
    best_score = float(scores[best_idx])

    log.info(
        f"  Embedding scores: {', '.join(f'[{i}]={float(scores[i]):.3f}' for i in range(len(results)))}")
    log.info(
        f"  Embedding best: [{best_idx}] {results[best_idx].get('title', '?')} (score={best_score:.3f}, threshold={threshold})")

    if best_score >= threshold:
        return best_idx, best_score

    return None, best_score


def score_all(
    title: str,
    author: str,
    results: list,
    model_name: str = "all-MiniLM-L6-v2",
    device: str = "",
) -> list[tuple[int, float]]:
    """
    Score all results, return list of (index, score) sorted best-first.
    Useful for diagnostics / debugging.
    """
    if not results:
        return []

    _load_model(model_name, device)

    if _model is None or _model is False:
        return []

    query = _build_text(title, author)
    candidates = [_result_text(r) for r in results]

    all_texts = [query] + candidates
    embeddings = _model.encode(
        all_texts, convert_to_tensor=True, show_progress_bar=False)

    from sentence_transformers.util import cos_sim
    scores = cos_sim(embeddings[0:1], embeddings[1:])[0]

    scored = [(i, float(scores[i])) for i in range(len(results))]
    scored.sort(key=lambda x: -x[1])
    return scored