mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-08-17 05:30:34 +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
81b6253261
commit
a3ff1f1444
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