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 ]