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)
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
#!/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()
|