File size: 3,735 Bytes
67a0f35 | 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 100 101 | """Memory-lean data prep using polars directly on the hub parquet shards."""
import glob
import os
import sys
import numpy as np
import polars as pl
N_TRAIN_Q = 80_000
N_DEV_Q = 2_000
N_NEGS = 4
SEED = 42
SNAP = glob.glob(os.path.expanduser(
"~/.cache/huggingface/hub/datasets--sentence-transformers--msmarco/snapshots/*"))[0]
print("scanning margin-mse shards...", flush=True)
df = (
pl.scan_parquet(f"{SNAP}/bert-ensemble-margin-mse/*.parquet")
.select(
pl.col("query_id").cast(pl.Int64),
pl.col("positive_id").cast(pl.Int64),
pl.col("negative_id").cast(pl.Int64),
pl.col("score").cast(pl.Float32),
)
.collect()
)
print("rows:", df.height, flush=True)
# canonical positive per query = first occurrence; then 4 distinct negatives for it
first_pos = df.group_by("query_id").agg(pl.col("positive_id").first().alias("pos"))
g = (
df.join(first_pos, on="query_id")
.filter((pl.col("positive_id") == pl.col("pos")) & (pl.col("negative_id") != pl.col("pos")))
.unique(subset=["query_id", "negative_id"], keep="first", maintain_order=True)
.group_by("query_id")
.agg(
pl.col("pos").first(),
pl.col("negative_id").head(N_NEGS).alias("negs"),
pl.col("score").head(N_NEGS).alias("margins"),
)
.filter(pl.col("negs").list.len() >= N_NEGS)
)
print("eligible queries:", g.height, flush=True)
g = g.sample(n=min(N_TRAIN_Q + N_DEV_Q, g.height), seed=SEED, shuffle=True)
print("sampled:", g.height, flush=True)
qids = g["query_id"].to_numpy()
pos = g["pos"].to_numpy()
negs = np.stack(g["negs"].to_numpy()) # [N, 4]
margins = np.stack(g["margins"].to_numpy()).astype(np.float32)
need_docs = np.unique(np.concatenate([pos, negs.reshape(-1)]))
print("unique docs:", len(need_docs), "queries:", len(qids), flush=True)
print("loading corpus/queries parquet...", flush=True)
corpus_files = glob.glob(f"{SNAP}/corpus/*.parquet")
queries_files = glob.glob(f"{SNAP}/queries/*.parquet")
assert corpus_files and queries_files, "run: hf download for corpus/queries first"
corpus = (
pl.scan_parquet(corpus_files)
.select(pl.col("passage_id").cast(pl.Int64).alias("id"), pl.col("passage").alias("text"))
.filter(pl.col("id").is_in(need_docs))
.collect()
)
queries = (
pl.scan_parquet(queries_files)
.select(pl.col("query_id").cast(pl.Int64).alias("id"), pl.col("query").alias("text"))
.filter(pl.col("id").is_in(qids))
.collect()
)
print("resolved docs:", corpus.height, "queries:", queries.height, flush=True)
assert corpus.height == len(need_docs), "missing corpus texts"
assert queries.height == len(qids), "missing query texts"
qtext = dict(zip(queries["id"].to_list(), queries["text"].to_list()))
dtext = dict(zip(corpus["id"].to_list(), corpus["text"].to_list()))
clean = lambda s: " ".join(s.split())
id2idx = {}
with open("/home/anon/pog/data/manifest.tsv", "w") as f:
for q in qids:
id2idx[f"q{q}"] = len(id2idx)
f.write(f"q{q}\tq\t{clean(qtext[q])}\n")
for d in need_docs:
id2idx[f"d{d}"] = len(id2idx)
f.write(f"d{d}\td\t{clean(dtext[d])}\n")
print("manifest rows:", len(id2idx), flush=True)
q_idx = np.array([id2idx[f"q{q}"] for q in qids], dtype=np.int64)
p_idx = np.array([id2idx[f"d{d}"] for d in pos], dtype=np.int64)
n_idx = np.array([[id2idx[f"d{d}"] for d in row] for row in negs], dtype=np.int64)
tr = slice(0, N_TRAIN_Q)
dv = slice(N_TRAIN_Q, None)
np.savez("/home/anon/pog/data/train.npz", q=q_idx[tr], pos=p_idx[tr], neg=n_idx[tr], margin=margins[tr])
np.savez("/home/anon/pog/data/dev.npz", q=q_idx[dv], pos=p_idx[dv], neg=n_idx[dv], margin=margins[dv])
print(f"DONE train={N_TRAIN_Q} dev={len(qids)-N_TRAIN_Q}", flush=True)
|