semantic_search.py

import datasets
import numpy as np
import scipy.spatial
import scipy.special
import spaces
from sentence_transformers import CrossEncoder, SentenceTransformer

from table import BASE_REPO_ID

ds = datasets.load_dataset(BASE_REPO_ID, split="train")
ds = ds.rename_column("submission_number", "paper_id")
ds.add_faiss_index(column="embedding")

model = SentenceTransformer("all-MiniLM-L6-v2")
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")


@spaces.GPU(duration=5)
def semantic_search(query: str, candidate_pool_size: int = 100, score_threshold: float = 0.7) -> list[int]:
    query_vec = model.encode(query)
    _, retrieved_data = ds.get_nearest_examples("embedding", query_vec, k=candidate_pool_size)

    rerank_inputs = [
        [query, f"{title}\n{abstract}"]
        for title, abstract in zip(retrieved_data["title"], retrieved_data["abstract"], strict=True)
    ]
    rerank_scores = reranker.predict(rerank_inputs)
    sorted_indices = np.argsort(rerank_scores)[::-1]

    return [
        retrieved_data["paper_id"][i]
        for i in sorted_indices
        if scipy.special.expit(rerank_scores[i]) >= score_threshold
    ]