mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
allaun
88db8987b0
- 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)
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
#!/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()
|