Building a Semantic Search Engine and Open-Status Classifier over the ResearchMath-14k Dataset

from sentence_transformers import util
def search(query, k=5):
q = model.encode([query], normalize_embeddings=True)
sims = util.cos_sim(q, emb)[0].cpu().numpy()
idx = sims.argsort()[::-1][:k]
print(f’\n=== Query: “{query}” ===’)
for rank, i in enumerate(idx, 1):
row = work.iloc[i]
print(f”\n[{rank}] sim={sims[i]:.3f} | {row[‘taxonomy_level_1’]} ”
f”| status={row[‘open_status’]}”)
print(” “, row[TEXT_COL][:260].replace(“\n”, ” “), “…”)
search(“rational points on hyperelliptic curves”)
search(“multiplicativity of maximal output p-norm of a quantum channel”)
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, ConfusionMatrixDisplay
y = work[“open_status”].values
Xtr, Xte, ytr, yte = train_test_split(
emb, y, test_size=0.25, random_state=RANDOM_STATE, stratify=y)
clf = LogisticRegression(max_iter=2000, class_weight=”balanced”, C=2.0)
clf.fit(Xtr, ytr)
pred = clf.predict(Xte)
print(“\n=== open_status classifier (embeddings + logistic regression) ===”)
print(classification_report(yte, pred))
fig, ax = plt.subplots(figsize=(7, 6))
ConfusionMatrixDisplay.from_predictions(
yte, pred, ax=ax, cmap=”Blues”, xticks_rotation=45,
normalize=”true”, values_format=”.2f”)
ax.set_title(“open_status confusion matrix (row-normalized)”)
plt.tight_layout(); plt.show()
sims = util.cos_sim(emb, emb).cpu().numpy()
np.fill_diagonal(sims, 0)
i, j = np.unravel_index(sims.argmax(), sims.shape)
print(f”\nMost similar pair (cos={sims[i, j]:.3f}):”)
for n in (i, j):
print(f”\n paper_id={work.iloc[n][‘paper_id’]} | ”
f”{work.iloc[n][‘taxonomy_level_1’]}”)
print(” “, work.iloc[n][TEXT_COL][:240].replace(“\n”, ” “), “…”)
print(“\nDone. Set SAMPLE_SIZE=None at the top to run on the full 14.1k rows.”)