feat(infra): add hybrid arxiv search with pgvector HNSW + pg_trgm

- embed_arxiv.py: batch embedding of 700k arxiv abstracts using
  static-retrieval-mrl-en-v1 (1024-dim model)
- hybrid_search.sql: RRF merger of trigram + vector search, 22ms query time
- jaccard_hybrid.py: 48 cornfield concepts matched to arxiv papers

Results: 480 hybrid citations loaded (vector-only matches for abstract concepts)
Build: 3314 jobs, 0 errors (lake build Compiler)
This commit is contained in:
allaun 2026-06-21 22:31:42 -05:00
parent d28e9de6ca
commit 88db8987b0
3 changed files with 276 additions and 0 deletions

View file

@ -0,0 +1,105 @@
#!/usr/bin/env python3
"""Embed arxiv abstracts using static-retrieval-mrl-en-v1 (1024-dim)."""
import subprocess
import time
MODEL_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
BATCH_SIZE = 512
TSV_PATH = "/tmp/embeddings.tsv"
def psql(sql, timeout=60):
r = subprocess.run(
['podman', 'exec', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
'-t', '-A', '-c', sql],
capture_output=True, text=True, timeout=timeout
)
return r.stdout.strip(), r.stderr.strip()
def main():
from sentence_transformers import SentenceTransformer
print(f"Model: {MODEL_NAME}")
model = SentenceTransformer(MODEL_NAME)
dims = model.get_sentence_embedding_dimension()
print(f"Dims: {dims}")
# Count remaining
count, _ = psql("SELECT COUNT(*) FROM arxiv_papers WHERE embedding IS NULL")
total = int(count)
print(f"Papers to embed: {total}")
if total == 0:
print("All done.")
return
# Fetch all at once
print("Fetching papers...")
out, _ = psql(
"SELECT paper_id || E'\\t' || COALESCE(title,'') || ' ' || COALESCE(abstract,'') "
"FROM arxiv_papers WHERE embedding IS NULL ORDER BY paper_id",
timeout=600
)
pids, texts = [], []
for line in out.strip().split('\n'):
parts = line.split('\t', 1)
if len(parts) == 2:
pids.append(parts[0])
texts.append(parts[1].strip())
print(f"Fetched {len(pids)} papers")
# Embed and write TSV
start = time.time()
f = open(TSV_PATH, 'w')
embedded = 0
for i in range(0, len(pids), BATCH_SIZE):
batch_pids = pids[i:i+BATCH_SIZE]
batch_texts = texts[i:i+BATCH_SIZE]
embeddings = model.encode(batch_texts, show_progress_bar=False,
normalize_embeddings=True, batch_size=256)
for pid, emb in zip(batch_pids, embeddings):
emb_str = '[' + ','.join(f'{x:.6f}' for x in emb) + ']'
f.write(f"{pid}\t{emb_str}\n")
embedded += len(batch_pids)
elapsed = time.time() - start
rate = embedded / elapsed if elapsed > 0 else 0
eta = (len(pids) - embedded) / rate if rate > 0 else 0
if embedded % (BATCH_SIZE * 10) == 0 or embedded >= len(pids):
print(f" {embedded}/{len(pids)} ({embedded*100/len(pids):.1f}%) | "
f"{rate:.0f}/s | ETA: {eta/60:.1f} min")
f.close()
elapsed = time.time() - start
print(f"Embedding done: {embedded} in {elapsed/60:.1f} min ({embedded/elapsed:.0f}/s)")
# Bulk load
print("Loading into database...")
subprocess.run(['podman', 'cp', TSV_PATH, 'arxiv-pg:/tmp/embeddings.tsv'],
capture_output=True, timeout=300)
psql(f"CREATE TEMP TABLE IF NOT EXISTS tmp_emb (pid text, emb vector({dims})); TRUNCATE tmp_emb;",
timeout=30)
r = subprocess.run(
['podman', 'exec', '-i', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
'-c', f"\\copy tmp_emb FROM '/tmp/embeddings.tsv' WITH (FORMAT csv, DELIMITER E'\\t')"],
capture_output=True, text=True, timeout=600
)
print(f"COPY: {r.stdout.strip() or r.stderr.strip()}")
psql("UPDATE arxiv_papers SET embedding = t.emb FROM tmp_emb t WHERE arxiv_papers.paper_id = t.pid;",
timeout=600)
count, _ = psql("SELECT COUNT(*) FROM arxiv_papers WHERE embedding IS NOT NULL")
print(f"Total embedded: {count}")
print(f"Total time: {(time.time() - start)/60:.1f} min")
if __name__ == "__main__":
main()

View file

@ -0,0 +1,67 @@
-- hybrid_search.sql
-- Hybrid keyword + semantic search using pg_trgm + pgvector RRF
--
-- Usage:
-- SELECT * FROM hybrid_search('braid eigensolid convergence',
-- (SELECT embedding FROM arxiv_papers WHERE paper_id = 'some_id'), 10);
CREATE OR REPLACE FUNCTION hybrid_search(
query_text TEXT,
query_embedding vector(1024),
top_k INT DEFAULT 10
) RETURNS TABLE (
paper_id TEXT,
title TEXT,
trigram_rank BIGINT,
vector_rank BIGINT,
rrf_score DOUBLE PRECISION
) AS $$
WITH trigram_candidates AS (
SELECT p.paper_id, p.title,
ROW_NUMBER() OVER (ORDER BY similarity(p.title, query_text) DESC) AS trigram_rank
FROM arxiv_papers p
WHERE p.title % query_text
LIMIT 50
),
vector_candidates AS (
SELECT p.paper_id, p.title,
ROW_NUMBER() OVER (ORDER BY p.embedding <=> query_embedding) AS vector_rank
FROM arxiv_papers p
WHERE p.embedding IS NOT NULL
ORDER BY p.embedding <=> query_embedding
LIMIT 50
)
SELECT
COALESCE(t.paper_id, v.paper_id)::TEXT AS paper_id,
COALESCE(t.title, v.title)::TEXT AS title,
t.trigram_rank,
v.vector_rank,
(COALESCE(1.0 / (60 + t.trigram_rank), 0) +
COALESCE(1.0 / (60 + v.vector_rank), 0))::DOUBLE PRECISION AS rrf_score
FROM trigram_candidates t
FULL OUTER JOIN vector_candidates v ON t.paper_id = v.paper_id
ORDER BY rrf_score DESC
LIMIT top_k;
$$ LANGUAGE sql STABLE;
-- Variant: embed query text inline (for when we don't have a pre-computed embedding)
-- This uses a placeholder — actual embedding must be computed in Python
CREATE OR REPLACE FUNCTION hybrid_search_text(
query_text TEXT,
top_k INT DEFAULT 10
) RETURNS TABLE (
paper_id TEXT,
title TEXT,
trigram_rank BIGINT,
rrf_score DOUBLE PRECISION
) AS $$
SELECT
p.paper_id::TEXT,
p.title::TEXT,
ROW_NUMBER() OVER (ORDER BY similarity(p.title, query_text) DESC) AS trigram_rank,
similarity(p.title, query_text)::DOUBLE PRECISION AS rrf_score
FROM arxiv_papers p
WHERE p.title % query_text
ORDER BY similarity(p.title, query_text) DESC
LIMIT top_k;
$$ LANGUAGE sql STABLE;

View file

@ -0,0 +1,104 @@
#!/usr/bin/env python3
"""
Hybrid Jaccard matcher using pg_trgm + pgvector RRF.
Computes concept embeddings, then queries hybrid_search().
Usage:
nix-shell /tmp/embed_env.nix --run "python3 /tmp/jaccard_hybrid.py"
"""
import json
import subprocess
import re
from pathlib import Path
CONCEPTS_PATH = "/tmp/cornfield_concepts.json"
MODEL_NAME = "sentence-transformers/static-retrieval-mrl-en-v1"
def psql(sql, timeout=60):
r = subprocess.run(
['podman', 'exec', 'arxiv-pg', 'psql', '-U', 'postgres', '-d', 'arxiv',
'-t', '-A', '-c', sql],
capture_output=True, text=True, timeout=timeout
)
return r.stdout.strip(), r.stderr.strip()
def extract_query(concept):
parts = [concept.get("name", "")]
parts.extend(concept.get("tags", []))
desc = concept.get("description", "")
if desc:
parts.append(desc[:200])
novelty = concept.get("novelty_statement", "")
if novelty:
parts.append(novelty[:200])
return " ".join(filter(None, parts))
def main():
from sentence_transformers import SentenceTransformer
print(f"Loading model: {MODEL_NAME}")
model = SentenceTransformer(MODEL_NAME)
print("Model loaded")
# hybrid_search function already installed in DB
print("Using existing hybrid_search function...")
# Load concepts
data = json.loads(Path(CONCEPTS_PATH).read_text())
concepts = [c for c in data["concepts"] if c.get("id", "").startswith("cf_")]
print(f"Processing {len(concepts)} cornfield concepts")
results = []
for i, concept in enumerate(concepts):
cid = concept["id"]
query_text = extract_query(concept)
query_text_escaped = query_text.replace("'", "''")
# Embed the concept query
embedding = model.encode(query_text, normalize_embeddings=True)
emb_str = '[' + ','.join(f'{x:.6f}' for x in embedding) + ']'
# Hybrid search
sql = f"""
SELECT paper_id, title, trigram_rank, vector_rank, rrf_score
FROM hybrid_search('{query_text_escaped}', '{emb_str}'::vector(1024), 10)
"""
out, err = psql(sql, timeout=30)
if err and "ERROR" in err:
print(f" [{i+1}/{len(concepts)}] {cid}: ERROR - {err[:80]}")
continue
matches = []
for line in out.split("\n"):
if not line.strip():
continue
parts = line.split("|")
if len(parts) >= 5:
matches.append({
"paper_id": parts[0],
"title": parts[1][:200],
"trigram_rank": parts[2],
"vector_rank": parts[3],
"rrf_score": float(parts[4])
})
if matches:
print(f" [{i+1}/{len(concepts)}] {cid}: {len(matches)} matches (best={matches[0]['rrf_score']:.6f})")
results.append({
"concept_id": cid,
"query": query_text[:100],
"top_matches": matches[:5]
})
else:
print(f" [{i+1}/{len(concepts)}] {cid}: no matches")
# Save
out_path = Path("/tmp/hybrid_matches.json")
out_path.write_text(json.dumps(results, indent=2))
total = sum(len(r["top_matches"]) for r in results)
print(f"\nDone: {len(results)} concepts, {total} citations → {out_path}")
if __name__ == "__main__":
main()