backend/astroparse_api/retrieve.py

import os
import re
from pathlib import Path

import numpy as np

_POS_TAGS = {"PROPN", "ADJ", "NOUN"}
_nlp = None


def _get_nlp():
    global _nlp
    if _nlp is None:
        import spacy
        _nlp = spacy.load("en_core_web_sm")
    return _nlp


def get_keywords(text: str) -> list[str]:
    doc = _get_nlp()(text.lower())
    return [t.text for t in doc if t.pos_ in _POS_TAGS and not t.is_stop and not t.is_punct]


def keyword_weights(query_kws: list[str], candidate_kws: list[list[str]]) -> np.ndarray:
    """Pathfinder scheme (pathfinder_setup_fns.RetrievalSystem.rank_and_filter)."""
    w = np.full(len(candidate_kws), 0.1)
    for qk in query_kws:
        qk = qk.lower()
        for i, kws in enumerate(candidate_kws):
            for k in kws:
                if qk in k.lower():
                    w[i] += 0.1
    return w / w.max()


# ---------------------------------------------------------------------------
# Corpus wrapper and per-paragraph retrieval
# ---------------------------------------------------------------------------

from .schemas import LitPaper, Scores, Paragraph  # noqa: E402

EMBED_MODEL = "text-embedding-3-small"
_ARXIV_ID = re.compile(r"^\d{4}\.\d{4,5}$")


# Compact local corpus layout (built by scripts/build_corpus.py from the HF
# parquet shards; the full dataset is ~20 GB with float64 embeds and citation
# graphs, far beyond this machine's disk/RAM budget):
#   - text columns only (title/abstract/authors/date/keywords/bibcode/arxiv_id)
#     as parquet shards in ASTROPARSE_CORPUS_DIR
#   - a float16-quantized faiss index (IndexScalarQuantizer, METRIC_L2) at
#     ASTROPARSE_FAISS_PATH, memory-mapped at load so resident RAM stays low
_DEFAULT_DATA_DIR = Path(__file__).resolve().parent.parent / "data"


class Corpus:
    def __init__(self, dataset, index):
        self.ds = dataset
        self.index = index

    @property
    def arxiv_ids(self) -> set[str]:
        """Cached set of all arXiv ids in the corpus."""
        if not hasattr(self, "_arxiv_ids"):
            col = self.ds["arxiv_id"] if "arxiv_id" in self.ds.column_names else []
            self._arxiv_ids = {x for x in col if x}
        return self._arxiv_ids

    @property
    def bibcodes(self) -> set[str]:
        """Cached set of all bibcodes in the corpus."""
        if not hasattr(self, "_bibcodes"):
            col = self.ds["bibcode"] if "bibcode" in self.ds.column_names else []
            self._bibcodes = {x for x in col if x}
        return self._bibcodes

    @classmethod
    def load(cls) -> "Corpus":
        import faiss
        from datasets import load_dataset

        corpus_dir = Path(os.environ.get("ASTROPARSE_CORPUS_DIR", _DEFAULT_DATA_DIR / "corpus"))
        faiss_path = Path(os.environ.get("ASTROPARSE_FAISS_PATH", _DEFAULT_DATA_DIR / "astroparse_fp16.faiss"))
        if not corpus_dir.is_dir() or not faiss_path.is_file():
            raise RuntimeError(
                f"compact corpus not found ({corpus_dir}, {faiss_path}) — "
                "run: uv run python scripts/build_corpus.py"
            )
        ds = load_dataset("parquet", data_files=str(corpus_dir / "*.parquet"), split="train")
        index = faiss.read_index(str(faiss_path), faiss.IO_FLAG_MMAP)
        return cls(ds, index)

    def search(self, query_embedding: np.ndarray, k: int = 1000):
        q = np.asarray(query_embedding, dtype=np.float32).reshape(1, -1)
        distances, indices = self.index.search(q, k)
        # distance -> similarity (pathfinder convention); epsilon guards the
        # zero-distance case (a paragraph can be a corpus paper's own abstract)
        sims = 1.0 / (np.asarray(distances[0]) + 1e-10)
        return [int(i) for i in indices[0]], sims


def embed_batch(texts: list[str], api_key: str) -> list[np.ndarray]:
    from openai import OpenAI

    client = OpenAI(api_key=api_key)
    out: list[np.ndarray] = []
    for i in range(0, len(texts), 100):
        data = client.embeddings.create(input=texts[i : i + 100], model=EMBED_MODEL).data
        out += [np.array(d.embedding, dtype=np.float32) for d in data]
    return out


def _short_name(authors, year: int) -> str:
    first = authors[0] if isinstance(authors, list) else str(authors).split(";")[0].split(" and ")[0]
    last = first.split(",")[0].strip().split()[-1]
    return f"{last}{str(year)[2:]}"


def _lit_from_row(row: dict, embed_score: float, kw_score: float, rank: int) -> LitPaper:
    date = row["date"]
    if hasattr(date, "year"):
        year = date.year
    else:
        year = int(str(date)[:4])
    authors = ", ".join(row["authors"]) if isinstance(row["authors"], list) else str(row["authors"])
    # Real corpus: arxiv_id is a dedicated column; ads_id is a numeric ADS internal id.
    # Use arxiv_id directly but still validate format so old-style ids (astro-ph9804...) get suppressed.
    arxiv_id = str(row.get("arxiv_id", ""))
    arxiv = arxiv_id if _ARXIV_ID.match(arxiv_id) else ""
    return LitPaper(
        id=row["bibcode"],
        short=_short_name(row["authors"], year),
        title=row["title"],
        authors=authors,
        year=year,
        journal="",
        bibcode=row["bibcode"],
        arxiv=arxiv,
        abstract=row["abstract"],
        scores=Scores(embed=round(float(embed_score), 4), keywords=round(float(kw_score), 4), rank=rank),
    )


def retrieve_for_paragraphs(
    corpus: Corpus,
    paragraphs: list[Paragraph],
    embeddings: list[np.ndarray],
    top_k: int = 10,
    pool: int = 1000,
):
    lit_papers: dict[str, LitPaper] = {}
    lit_by_para: dict[str, list[str]] = {}
    for para, emb in zip(paragraphs, embeddings):
        indices, sims = corpus.search(np.asarray(emb, dtype=np.float32), k=pool)
        rows = corpus.ds[indices]
        kw = keyword_weights(get_keywords(para.text), rows["keywords"])
        embed_norm = sims / sims.max()
        combined = embed_norm * kw
        order = np.argsort(combined)[::-1][:top_k]
        ids = []
        for rank, j in enumerate(order, start=1):
            row = {k: rows[k][j] for k in rows}
            lit = _lit_from_row(row, embed_norm[j], kw[j], rank)
            if lit.id not in lit_papers:
                lit_papers[lit.id] = lit
            ids.append(lit.id)
        lit_by_para[para.id] = ids
    return lit_papers, lit_by_para