File size: 3,691 Bytes
e79d917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Unified, resumable embedding stage — replaces model-2's embed_train.py + embed_test.py.
python src/embed.py --split train --part pool --gpu 0
python src/embed.py --split test --part pool --gpu 1
python src/embed.py --split test --part s1 --gpu 1
Parts
-----
pool : Source 2 + Source 3 (S2 rows first, then S3) -- the search corpus, doc prefix
s1 : Source 1 -- the queries, query prefix
all : pool then s1
Notes
-----
* `--split train --part s1` is OPTIONAL and off by default in scripts/queue.sh.
run_train_eval.py only queries n_s1_ce+n_s1_stack+n_s1_val (=200k) of the 2.21M
train S1, and caches those itself. Embedding all 2.21M is the dead work that cost
model-2 ~50 min. `--split test --part s1` IS required (every test S1 must be scored).
* Every run is resumable and fingerprint-checked -- see encoder.py.
* --gpu is mandatory-by-default so nothing silently lands on a busy cuda:0.
"""
from __future__ import annotations
import argparse
import os
import sys
import time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pandas as pd # noqa: E402
from config import CFG, EMB, PATHS, emb_dir, model_spec # noqa: E402
from encoder import encode_to_shards, write_ids # noqa: E402
COLS = ["entity_id", "business_name", "business_address", "country"]
T0 = time.time()
def log(m):
print(f"[{time.time() - T0:7.0f}s] {m}", flush=True)
def record_text(name, addr, country):
"""Must stay byte-identical to the text used at search time."""
return f"{name} | {addr} | {country}"
def read_part(split, part):
p = PATHS[split]
if part == "pool":
df = pd.concat([pd.read_parquet(p["source2"], columns=COLS),
pd.read_parquet(p["source3"], columns=COLS)], ignore_index=True)
else:
df = pd.read_parquet(p["source1"], columns=COLS)
for c in COLS:
df[c] = df[c].fillna("").astype(str)
return df
def run(split, part, device, model_name=None):
hf_id, qpre, dpre, _dim, lic = model_spec(model_name)
prefix = dpre if part == "pool" else qpre
df = read_part(split, part)
log(f"{split}/{part}: {len(df):,} records | model={hf_id} ({lic}) | prefix={prefix!r}")
texts = [record_text(a, b, c) for a, b, c in
zip(df["business_name"], df["business_address"], df["country"])]
ids = df["entity_id"].tolist()
del df
d = emb_dir(f"{split}_{part}", model_name)
os.makedirs(d, exist_ok=True)
write_ids(os.path.join(d, "ids.txt"), ids)
return encode_to_shards(texts, ids, d, hf_id, prefix, EMB["max_len"], EMB["batch"],
device, EMB["shard_rows"], log)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--split", choices=["train", "test"], required=True)
ap.add_argument("--part", choices=["pool", "s1", "all"], default="pool")
ap.add_argument("--gpu", default=None,
help="GPU index, or 'cpu'. Required unless M3_GPU is set -- this is "
"deliberate: model-2 lost runs by defaulting to a saturated cuda:0.")
ap.add_argument("--model", default=None, help=f"override (default {EMB['model']})")
a = ap.parse_args()
gpu = a.gpu if a.gpu is not None else os.environ.get("M3_GPU")
if gpu is None:
sys.exit("refusing to guess a GPU: pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)")
device = "cpu" if str(gpu) == "cpu" else f"cuda:{int(gpu)}"
parts = ["pool", "s1"] if a.part == "all" else [a.part]
for p in parts:
run(a.split, p, device, a.model)
log("all requested parts complete")
if __name__ == "__main__":
main()
|