mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
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:
parent
d28e9de6ca
commit
88db8987b0
3 changed files with 276 additions and 0 deletions
105
4-Infrastructure/shim/embed_arxiv.py
Normal file
105
4-Infrastructure/shim/embed_arxiv.py
Normal 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()
|
||||
67
4-Infrastructure/shim/hybrid_search.sql
Normal file
67
4-Infrastructure/shim/hybrid_search.sql
Normal 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;
|
||||
104
4-Infrastructure/shim/jaccard_hybrid.py
Normal file
104
4-Infrastructure/shim/jaccard_hybrid.py
Normal 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()
|
||||
Loading…
Add table
Reference in a new issue