research: add reduced AttnRes runner
This commit is contained in:
wuyang
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Download and freeze the byte-level WikiText-2 corpus for the AttnRes study."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
|
||||
PROTOCOL_ID = "llm-atlas-k3-attnres-reduced-v1"
|
||||
DATASET_REPO = "Salesforce/wikitext"
|
||||
DATASET_REVISION = "b08601e04326c79dfdd32d625aee71d232d685c3"
|
||||
DATASET_VARIANT = "wikitext-2-raw-v1"
|
||||
SPLITS = ("train", "validation", "test")
|
||||
SEEDS = (2026073001, 2026073002, 2026073003)
|
||||
CONTEXT = 256
|
||||
FORMAL_STEPS = 2000
|
||||
FORMAL_BATCH = 32
|
||||
VALIDATION_WINDOWS = 64
|
||||
DIAGNOSTIC_WINDOWS = 16
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--cache-dir", type=Path, required=True)
|
||||
parser.add_argument("--manifest", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def bytes_sha256(payload: bytes) -> str:
|
||||
return hashlib.sha256(payload).hexdigest()
|
||||
|
||||
|
||||
def atomic_json(path: Path, value: dict[str, Any]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".tmp")
|
||||
temporary.write_text(
|
||||
json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
|
||||
)
|
||||
os.replace(temporary, path)
|
||||
|
||||
|
||||
def download(url: str, path: Path) -> None:
|
||||
if path.exists():
|
||||
return
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".part")
|
||||
request = urllib.request.Request(
|
||||
url,
|
||||
headers={"User-Agent": "llm-atlas-k3-attnres-reduced/1.0"},
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=120) as response:
|
||||
with temporary.open("wb") as output:
|
||||
while block := response.read(1024 * 1024):
|
||||
output.write(block)
|
||||
os.replace(temporary, path)
|
||||
|
||||
|
||||
def window_start(label: str, index: int, corpus_length: int, seed: int | None = None) -> int:
|
||||
fields = [PROTOCOL_ID, label]
|
||||
if seed is not None:
|
||||
fields.append(str(seed))
|
||||
fields.append(str(index))
|
||||
payload = "\0".join(fields).encode()
|
||||
value = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
|
||||
return value % (corpus_length - (CONTEXT + 1))
|
||||
|
||||
|
||||
def train_window_start(seed: int, step: int, row: int, corpus_length: int) -> int:
|
||||
payload = "\0".join(
|
||||
[PROTOCOL_ID, "train-window", str(seed), str(step), str(row)]
|
||||
).encode()
|
||||
value = int.from_bytes(hashlib.sha256(payload).digest()[:8], "big")
|
||||
return value % (corpus_length - (CONTEXT + 1))
|
||||
|
||||
|
||||
def concatenate_split(parquet_path: Path) -> tuple[bytes, int]:
|
||||
table = pq.read_table(parquet_path, columns=["text"])
|
||||
rows = table.column("text").to_pylist()
|
||||
payload = b"".join(((row or "") + "\n").encode("utf-8") for row in rows)
|
||||
return payload, len(rows)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
args.cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
split_manifest: dict[str, Any] = {}
|
||||
split_bytes: dict[str, bytes] = {}
|
||||
for split in SPLITS:
|
||||
relative = f"{DATASET_VARIANT}/{split}-00000-of-00001.parquet"
|
||||
url = (
|
||||
f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/"
|
||||
f"{DATASET_REVISION}/{relative}"
|
||||
)
|
||||
parquet_path = args.cache_dir / f"{split}.parquet"
|
||||
download(url, parquet_path)
|
||||
payload, rows = concatenate_split(parquet_path)
|
||||
binary_path = args.cache_dir / f"{split}.bin"
|
||||
if not binary_path.exists() or binary_path.read_bytes() != payload:
|
||||
temporary = binary_path.with_suffix(".bin.tmp")
|
||||
temporary.write_bytes(payload)
|
||||
os.replace(temporary, binary_path)
|
||||
split_bytes[split] = payload
|
||||
split_manifest[split] = {
|
||||
"source_path": relative,
|
||||
"source_url": url,
|
||||
"parquet_bytes": parquet_path.stat().st_size,
|
||||
"parquet_sha256": file_sha256(parquet_path),
|
||||
"rows": rows,
|
||||
"concatenated_bytes": len(payload),
|
||||
"concatenated_sha256": bytes_sha256(payload),
|
||||
"binary_path": str(binary_path),
|
||||
"binary_sha256": file_sha256(binary_path),
|
||||
}
|
||||
|
||||
train = split_bytes["train"]
|
||||
validation = split_bytes["validation"]
|
||||
|
||||
schedule_digest = hashlib.sha256()
|
||||
schedule_cells = 0
|
||||
for seed in SEEDS:
|
||||
for step in range(1, FORMAL_STEPS + 1):
|
||||
for row in range(FORMAL_BATCH):
|
||||
start = train_window_start(seed, step, row, len(train))
|
||||
schedule_digest.update(start.to_bytes(8, "big"))
|
||||
schedule_cells += 1
|
||||
|
||||
validation_starts = [
|
||||
window_start("validation-window", index, len(validation))
|
||||
for index in range(VALIDATION_WINDOWS)
|
||||
]
|
||||
diagnostic_starts = [
|
||||
window_start("diagnostic-window", index, len(validation))
|
||||
for index in range(DIAGNOSTIC_WINDOWS)
|
||||
]
|
||||
|
||||
def tensor_hash(starts: list[int]) -> str:
|
||||
digest = hashlib.sha256()
|
||||
for start in starts:
|
||||
digest.update(validation[start : start + CONTEXT + 1])
|
||||
return digest.hexdigest()
|
||||
|
||||
manifest = {
|
||||
"schema_version": 1,
|
||||
"protocol_id": PROTOCOL_ID,
|
||||
"status": "frozen-before-model-output",
|
||||
"dataset": {
|
||||
"repository": DATASET_REPO,
|
||||
"revision": DATASET_REVISION,
|
||||
"variant": DATASET_VARIANT,
|
||||
"preprocessing": (
|
||||
"parquet row order; (text or empty string) + LF; UTF-8; "
|
||||
"no normalization; vocabulary is raw bytes 0..255"
|
||||
),
|
||||
"splits": split_manifest,
|
||||
},
|
||||
"windows": {
|
||||
"context": CONTEXT,
|
||||
"target_bytes_per_window": CONTEXT,
|
||||
"seeds": list(SEEDS),
|
||||
"formal_steps": FORMAL_STEPS,
|
||||
"formal_batch": FORMAL_BATCH,
|
||||
"formal_schedule_cells": schedule_cells,
|
||||
"formal_schedule_sha256": schedule_digest.hexdigest(),
|
||||
"validation_starts": validation_starts,
|
||||
"validation_tensor_sha256": tensor_hash(validation_starts),
|
||||
"diagnostic_starts": diagnostic_starts,
|
||||
"diagnostic_tensor_sha256": tensor_hash(diagnostic_starts),
|
||||
},
|
||||
}
|
||||
atomic_json(args.manifest, manifest)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user