File size: 5,170 Bytes
dfe1826 | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Kaggle GPU job: embed the full Mathlib corpus and the queries of Q (Phase 1.3).
Runs on Kaggle, not locally. `scripts/kaggle_job.py kernel` copies this file into a
kernel folder, filling in the JOB block, and pushes it. The attached private dataset
(built by `scripts/kaggle_job.py dataset`) holds:
corpus.jsonl name, signature, text of every declaration in data/index/corpus.jsonl
queries.jsonl name, split, query of every item in Q (no gold sets)
models.yaml copy of config/models.yaml (model ids are never written in code)
encoders.py copy of src/formal_mas/retrieval/encoders.py
Output in /kaggle/working/<key>/: corpus.npy (fp16 storage, ℓ2-normalised), queries.json,
queries_no_instr.npy and queries_instr.npy (fp32 storage), meta.json.
With LIMIT set, only a seeded random sample of LIMIT corpus texts is embedded to measure speed.
"""
from __future__ import annotations
import hashlib
import importlib.util
import json
import random
import sys
import time
from pathlib import Path
# --- JOB (rewritten by scripts/kaggle_job.py) ---
KEYS = ["bge-base"]
LIMIT = None
PRECISION = "fp16"
BATCH_SIZE = 128
# --- end JOB ---
OUT = Path("/kaggle/working")
def find_input(name: str) -> Path:
hits = sorted(Path("/kaggle/input").rglob(name))
if not hits:
raise FileNotFoundError(name)
return hits[0]
def load_jsonl(path: Path) -> list[dict]:
with path.open(encoding="utf-8") as f:
return [json.loads(line) for line in f if line.strip()]
def names_sha(names: list[str]) -> str:
return hashlib.sha256("\n".join(names).encode("utf-8")).hexdigest()
def main() -> None:
import numpy as np
import sentence_transformers
import torch
import transformers
spec_mod = importlib.util.spec_from_file_location("encoders", find_input("encoders.py"))
encoders = importlib.util.module_from_spec(spec_mod)
sys.modules["encoders"] = encoders
spec_mod.loader.exec_module(encoders)
specs = encoders.load_specs(find_input("models.yaml"))
corpus = load_jsonl(find_input("corpus.jsonl"))
queries = load_jsonl(find_input("queries.jsonl"))
total = len(corpus)
if LIMIT:
corpus = random.Random(0).sample(corpus, LIMIT)
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "none"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"gpu={gpu} corpus={len(corpus)}/{total} queries={len(queries)}", flush=True)
for key in KEYS:
spec = specs[key]
out = OUT / key
out.mkdir(parents=True, exist_ok=True)
t0 = time.perf_counter()
enc = encoders.Encoder(spec, batch_size=BATCH_SIZE, device=device, precision=PRECISION)
load_s = time.perf_counter() - t0
enc.encode([spec.doc_text(r) for r in corpus[:BATCH_SIZE]]) # warm-up
if device == "cuda":
torch.cuda.synchronize()
t0 = time.perf_counter()
vecs = enc.encode([spec.doc_text(r) for r in corpus])
if device == "cuda":
torch.cuda.synchronize()
corpus_s = time.perf_counter() - t0
np.save(out / "corpus.npy", vecs.astype(np.float16))
qtexts = [r["query"] for r in queries]
np.save(out / "queries_no_instr.npy", enc.encode([spec.query_text(t, False) for t in qtexts]))
if spec.query_instruction:
np.save(out / "queries_instr.npy", enc.encode([spec.query_text(t, True) for t in qtexts]))
(out / "queries.json").write_text(json.dumps([r["name"] for r in queries]), encoding="utf-8")
rate = len(corpus) / corpus_s
meta = {
"embedding_key": key,
"hf_id": spec.hf_id,
"hf_revision": encoders.model_revision(spec.hf_id),
"backend": spec.backend,
"max_length": spec.max_length,
"doc_prefix": spec.doc_prefix,
"doc_template": spec.doc_template,
"query_instruction": spec.query_instruction,
"dim": enc.dim,
"precision": PRECISION,
"storage": {"corpus": "float16", "queries": "float32"},
"batch_size": BATCH_SIZE,
"n_corpus": len(corpus),
"corpus_names_sha256": names_sha([r["name"] for r in corpus]),
"limit": LIMIT,
"n_queries": len(queries),
"device": gpu,
"versions": {"torch": torch.__version__, "transformers": transformers.__version__,
"sentence_transformers": sentence_transformers.__version__},
"model_load_s": round(load_s, 1),
"corpus_embed_s": round(corpus_s, 1),
"texts_per_s": round(rate, 1),
"est_full_corpus_min": round(total / rate / 60, 1) if LIMIT else None,
}
(out / "meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
print(json.dumps({k: meta[k] for k in ("embedding_key", "texts_per_s", "corpus_embed_s",
"est_full_corpus_min", "dim", "precision")}), flush=True)
del enc
if device == "cuda":
torch.cuda.empty_cache()
if __name__ == "__main__":
main()
|