scripts/recommender.py

"""
CrossPaper recommendation engine.

Retrieves candidate papers via FAISS nearest-neighbor search on fine-tuned
embeddings, then reranks using Maximal Marginal Relevance (MMR) to balance
relevance with disciplinary diversity.

Supports both the fine-tuned model (main recommendations) and the base model
(for before/after comparison in the demo).

AI Attribution: MMR implementation and diversity scoring logic
assisted by Claude (Anthropic).
"""

import pickle
from pathlib import Path

import faiss
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer


DATA_DIR = Path("data/processed")
BASE_MODEL_DIR = Path("models/base")
FINETUNED_MODEL_DIR = Path("models/fine_tuned")

# Subtracted from a candidate's MMR score when its field already appears in the
# selected set. Sized against the score's own range: relevance and max_sim are
# cosine similarities, so scores sit roughly in [-1, 1] and 0.15 is a real but
# not overwhelming push toward an unrepresented field.
FIELD_REPEAT_PENALTY = 0.15


class CrossPaperRecommender:
    """Recommends cross-disciplinary papers with diversity-aware reranking.

    Uses FAISS for fast retrieval and MMR for ensuring recommendations
    span multiple fields rather than clustering in one field.
    """

    def __init__(self, data_dir=DATA_DIR, model_dir=FINETUNED_MODEL_DIR):
        """Initialize the recommender.

        Args:
            data_dir: Directory containing FAISS indexes and paper metadata.
            model_dir: Directory containing the sentence-transformer model.
        """
        self.data_dir = Path(data_dir)
        self.model_dir = Path(model_dir)
        self.model = None
        self.index = None
        self.metadata = None
        self.embeddings = None

    def load(self, index_name="finetuned"):
        """Load model, index, and metadata into memory.

        Args:
            index_name: Which index to load ('base' or 'finetuned').
        """
        model_path = (
            BASE_MODEL_DIR if index_name == "base" else FINETUNED_MODEL_DIR
        )
        print(f"Loading {index_name} model from {model_path}...")

        # device="cpu" is required, not a preference.
        #
        # sentence-transformers picks its device from torch.cuda.is_available()
        # when none is given. The ZeroGPU runtime patches that to return True so
        # apps believe a GPU exists, but the real device is only attached inside
        # a @spaces.GPU call. A model loaded outside one lands on a device that
        # is never materialised, and encode() then returns zero vectors — which
        # fails silently, because FAISS still returns k results ranked by index
        # order rather than similarity.
        #
        # This model is 22M parameters and retrieval is a CPU FAISS index, so
        # there is nothing to gain from a GPU here anyway.
        self.model = SentenceTransformer(str(model_path), device="cpu")

        index_path = self.data_dir / f"{index_name}.index"
        print(f"Loading FAISS index from {index_path}...")
        self.index = faiss.read_index(str(index_path))

        embeddings_path = self.data_dir / f"{index_name}_embeddings.npy"
        self.embeddings = np.load(str(embeddings_path))

        metadata_path = self.data_dir / "paper_metadata.pkl"
        self.metadata = pd.read_pickle(str(metadata_path))

        print(f"  Ready: {self.index.ntotal} papers indexed")

    def retrieve(self, query, top_k=50):
        """Retrieve top-k candidate papers by embedding similarity.

        Args:
            query: Natural language query string.
            top_k: Number of candidates to retrieve (before reranking).

        Returns:
            Tuple of (similarity scores array, candidate indices array).
        """
        query_embedding = self.model.encode(
            [query], normalize_embeddings=True
        ).astype(np.float32)

        scores, indices = self.index.search(query_embedding, top_k)

        # A degenerate query vector fails silently: FAISS still returns k
        # results, every inner product is zero, and index order decides the
        # ranking. The corpus is ordered by field, so a broken encoder returns
        # the head of the index and looks exactly like a model that only knows
        # one field. This check makes that failure loud.
        norm = float(np.linalg.norm(query_embedding))
        healthy = (
            np.isfinite(query_embedding).all()
            and abs(norm - 1.0) < 0.01
            and indices[0][0] >= 5
        )
        if healthy:
            print(
                f"[encoder-check] ok norm={norm:.4f} "
                f"top_idx={indices[0][:3].tolist()} "
                f"top_scores={[round(float(s), 3) for s in scores[0][:3]]}",
                flush=True,
            )
        else:
            print(
                f"[encoder-check] FAILED norm={norm:.4f} "
                f"finite={bool(np.isfinite(query_embedding).all())} "
                f"emb_head={[round(float(v), 4) for v in query_embedding[0][:5]]} "
                f"idx_head={indices[0][:5].tolist()} "
                f"score_head={[round(float(s), 4) for s in scores[0][:5]]}",
                flush=True,
            )

        return scores[0], indices[0]

    def mmr_rerank(self, query, candidates_idx, candidates_scores, top_n=10, lambda_param=0.6):
        """Rerank candidates using Maximal Marginal Relevance.

        Balances relevance (similarity to query) with diversity (dissimilarity
        to already-selected papers), with an additional field diversity
        bonus for papers from underrepresented fields.

        MMR(d) = lambda * sim(q, d) - (1 - lambda) * max(sim(d, d_j) for d_j in selected)

        An additional field penalty is applied: if a paper's field
        already appears in the selected set, its MMR score is reduced. This
        encourages the final list to span multiple fields.

        Args:
            query: Original query string (unused, scores pre-computed).
            candidates_idx: Array of candidate paper indices.
            candidates_scores: Array of similarity scores for candidates.
            top_n: Number of papers to return after reranking.
            lambda_param: Relevance vs. diversity tradeoff (0=pure diversity, 1=pure relevance).

        Returns:
            List of dictionaries with paper info and scores.
        """
        selected = []
        selected_indices = []
        remaining = list(range(len(candidates_idx)))

        for _ in range(min(top_n, len(candidates_idx))):
            best_score = -float("inf")
            best_idx = -1

            for i in remaining:
                paper_idx = candidates_idx[i]
                relevance = candidates_scores[i]

                # Diversity: max similarity to any already-selected paper
                if selected_indices:
                    candidate_emb = self.embeddings[paper_idx].reshape(1, -1)
                    selected_embs = self.embeddings[selected_indices]
                    similarities = np.dot(selected_embs, candidate_emb.T).flatten()
                    max_sim = np.max(similarities)
                else:
                    max_sim = 0.0

                mmr_score = lambda_param * relevance - (1 - lambda_param) * max_sim

                # Field diversity penalty, applied additively.
                #
                # A multiplicative penalty inverts: mmr_score is negative
                # whenever the diversity term dominates (low lambda), and
                # scaling a negative number down by a factor raises it, which
                # turns the penalty into a reward for repeating a field.
                # Subtracting a constant keeps the direction stable at every
                # lambda and on both sides of zero.
                paper_field = self.metadata.iloc[paper_idx]["field"]
                selected_fields = [
                    self.metadata.iloc[idx]["field"] for idx in selected_indices
                ]
                if paper_field in selected_fields:
                    mmr_score -= FIELD_REPEAT_PENALTY

                if mmr_score > best_score:
                    best_score = mmr_score
                    best_idx = i

            if best_idx == -1:
                break

            paper_idx = candidates_idx[best_idx]
            selected_indices.append(paper_idx)

            paper_row = self.metadata.iloc[paper_idx]
            selected.append({
                "title": paper_row["title"],
                "abstract": paper_row.get("abstract", "")[:300],
                "field": paper_row["field"],
                "year": int(paper_row.get("year", 0)),
                "cited_by_count": int(paper_row.get("cited_by_count", 0)),
                "relevance_score": float(candidates_scores[best_idx]),
                "mmr_score": float(best_score),
            })

            remaining.remove(best_idx)

        return selected

    def recommend(self, query, top_n=10, lambda_param=0.6):
        """Generate recommendations for a query with diversity reranking.

        This is the main entry point. Retrieves candidates via FAISS,
        then applies MMR reranking to balance relevance with field
        diversity.

        Args:
            query: Natural language description of research interest.
            top_n: Number of recommendations to return.
            lambda_param: Relevance vs. diversity tradeoff.

        Returns:
            Dictionary with recommendations list and diversity metrics.
        """
        scores, indices = self.retrieve(query, top_k=top_n * 5)
        recommendations = self.mmr_rerank(
            query, indices, scores, top_n=top_n, lambda_param=lambda_param
        )

        diversity_metrics = self._compute_diversity(recommendations)

        return {
            "recommendations": recommendations,
            "diversity": diversity_metrics,
        }

    def _compute_diversity(self, recommendations):
        """Compute field diversity metrics for a recommendation set.

        Args:
            recommendations: List of recommendation dictionaries.

        Returns:
            Dictionary with diversity metrics including Shannon entropy,
            field distribution, and cross-field hit rate.
        """
        if not recommendations:
            return {"entropy": 0.0, "distribution": {}, "cross_field_rate": 0.0}

        fields = [r["field"] for r in recommendations]
        unique, counts = np.unique(fields, return_counts=True)
        probs = counts / counts.sum()

        # Shannon entropy (higher = more diverse)
        entropy = -np.sum(probs * np.log2(probs + 1e-10))

        # Distribution as percentages
        distribution = {
            disc: float(count / len(fields))
            for disc, count in zip(unique, counts)
        }

        # Cross-field rate (fraction of results NOT from the dominant field)
        dominant_fraction = max(probs)
        cross_rate = 1.0 - dominant_fraction

        return {
            "entropy": float(entropy),
            "distribution": distribution,
            "cross_field_rate": float(cross_rate),
            "num_fields": int(len(unique)),
        }


def main():
    """Quick smoke test for the recommender."""
    recommender = CrossPaperRecommender()
    recommender.load(index_name="finetuned")

    test_queries = [
        "attention mechanism in visual processing",
        "reinforcement learning for decision making",
        "gene expression regulation in neural development",
    ]

    for query in test_queries:
        print(f"\nQuery: {query}")
        print("-" * 60)
        result = recommender.recommend(query, top_n=5)

        for i, rec in enumerate(result["recommendations"], 1):
            print(f"  {i}. [{rec['field']}] {rec['title'][:80]}")
            print(f"     relevance={rec['relevance_score']:.3f}  mmr={rec['mmr_score']:.3f}")

        div = result["diversity"]
        print(f"  Diversity: entropy={div['entropy']:.2f}, "
              f"fields={div['num_fields']}, "
              f"cross_rate={div['cross_field_rate']:.0%}")


if __name__ == "__main__":
    main()