diff --git a/4-Infrastructure/shim/embed_arxiv.py b/4-Infrastructure/shim/embed_arxiv.py new file mode 100644 index 00000000..3dc60200 --- /dev/null +++ b/4-Infrastructure/shim/embed_arxiv.py @@ -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() diff --git a/4-Infrastructure/shim/hybrid_search.sql b/4-Infrastructure/shim/hybrid_search.sql new file mode 100644 index 00000000..cce18a0a --- /dev/null +++ b/4-Infrastructure/shim/hybrid_search.sql @@ -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; diff --git a/4-Infrastructure/shim/jaccard_hybrid.py b/4-Infrastructure/shim/jaccard_hybrid.py new file mode 100644 index 00000000..104c0d78 --- /dev/null +++ b/4-Infrastructure/shim/jaccard_hybrid.py @@ -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()