WP44: delta embedding mode (801,093 works added by corpus v3)
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
1 changed file +19 −0
modified
embeddings/embed_corpus.py
+19 −0
@@ -25,6 +25,22 @@ def prepare(root, shard_rows, title_only_fallback): | ||
| 25 | 25 | print("prepared", n, "texts in", len(glob.glob(out + "/shard=*")), "shards") |
| 26 | 26 | |
| 27 | 27 | |
| 28 | +def prepare_delta(root, shard_rows, model_name, start_shard=100): | |
| 29 | + """Shards for eligible works with an abstract that have no embedding yet (ids read from the existing shard_*.ids.npy).""" | |
| 30 | + import duckdb, glob as _g | |
| 31 | + tag = model_name.replace("/", "__"); done = np.concatenate([np.load(f) for f in _g.glob(os.path.join(root, "data/processed/embeddings", tag, "shard_*.ids.npy"))]) if _g.glob(os.path.join(root, "data/processed/embeddings", tag, "shard_*.ids.npy")) else np.array([], dtype=np.int64) | |
| 32 | + con = duckdb.connect(os.path.join(root, "data/processed/wp44.duckdb"), read_only=True); con.execute("SET threads TO 8;") | |
| 33 | + con.register("done_ids", __import__("pandas").DataFrame({"work_id": done})) | |
| 34 | + out = os.path.join(root, "data/processed/embeddings/input"); os.makedirs(out, exist_ok=True) | |
| 35 | + con.execute("""CREATE OR REPLACE TEMP TABLE emb_in AS | |
| 36 | + SELECT w.id AS work_id, w.title || '. ' || coalesce(substr(a.abstract, 1, 2000), '') AS text | |
| 37 | + FROM corpus_works w JOIN oa_abstracts a ON a.work_id = w.id LEFT JOIN done_ids d ON d.work_id = w.id | |
| 38 | + WHERE w.eligible AND d.work_id IS NULL ORDER BY w.id""") | |
| 39 | + n = con.execute("SELECT count(*) FROM emb_in").fetchone()[0] | |
| 40 | + con.execute(f"COPY (SELECT *, {start_shard} + (row_number() OVER (ORDER BY work_id) - 1) // {shard_rows} AS shard FROM emb_in) TO '{out}' (FORMAT PARQUET, PARTITION_BY (shard), OVERWRITE_OR_IGNORE)") | |
| 41 | + print("delta prepared", n, "texts from shard", start_shard, "to", start_shard + (n - 1) // shard_rows) | |
| 42 | + | |
| 43 | + | |
| 28 | 44 | def run(root, model_name, shard_from, shard_to, batch, max_len): |
| 29 | 45 | import torch, pyarrow.parquet as pq |
| 30 | 46 | from sentence_transformers import SentenceTransformer |
@@ -52,7 +68,10 @@ def main(): | ||
| 52 | 68 | ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true") |
| 53 | 69 | ap.add_argument("--model", default="allenai/specter2_base"); ap.add_argument("--shard-rows", type=int, default=50000); ap.add_argument("--shard-from", type=int, default=0) |
| 54 | 70 | ap.add_argument("--shard-to", type=int, default=10**6); ap.add_argument("--batch", type=int, default=128); ap.add_argument("--max-len", type=int, default=256); ap.add_argument("--title-only-fallback", action="store_true") |
| 71 | + ap.add_argument("--prepare-delta", action="store_true", help="shards (numbered from 100) for eligible works not yet embedded") | |
| 55 | 72 | a = ap.parse_args() |
| 73 | + if a.prepare_delta: | |
| 74 | + prepare_delta(a.root, a.shard_rows, a.model) | |
| 56 | 75 | if a.prepare: |
| 57 | 76 | prepare(a.root, a.shard_rows, a.title_only_fallback) |
| 58 | 77 | if a.run: |
| 59 | 78 | |