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 | print("prepared", n, "texts in", len(glob.glob(out + "/shard=*")), "shards") | 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 | def run(root, model_name, shard_from, shard_to, batch, max_len): | 44 | def run(root, model_name, shard_from, shard_to, batch, max_len): | ||
| 29 | import torch, pyarrow.parquet as pq | 45 | import torch, pyarrow.parquet as pq | ||
| 30 | from sentence_transformers import SentenceTransformer | 46 | from sentence_transformers import SentenceTransformer | ||
@@ -52,7 +68,10 @@ def main(): | |||||
| 52 | ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true") | 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 | 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) | 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 | 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") | 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 | a = ap.parse_args() | 72 | a = ap.parse_args() | ||
| 73 | + if a.prepare_delta: | ||||
| 74 | + prepare_delta(a.root, a.shard_rows, a.model) | ||||
| 56 | if a.prepare: | 75 | if a.prepare: | ||
| 57 | prepare(a.root, a.shard_rows, a.title_only_fallback) | 76 | prepare(a.root, a.shard_rows, a.title_only_fallback) | ||
| 58 | if a.run: | 77 | if a.run: | ||
| 59 | 78 | ||||