| import hashlib |
| import json |
| import os |
| import random |
| import tempfile |
| from pathlib import Path |
|
|
| import pyarrow.parquet as pq |
| import torch |
| import torch.distributed as dist |
| from filelock import FileLock |
|
|
|
|
| def parquet_files(data_dir, split): |
| if not data_dir: |
| raise ValueError("A local dataset directory is required.") |
| root = Path(data_dir).expanduser() |
| if not root.is_dir(): |
| raise FileNotFoundError(f"Local dataset directory does not exist: {root}") |
| files = sorted((root / "data").glob(f"{split}-*.parquet")) |
| if not files: |
| raise FileNotFoundError(f"No {split} parquet shards under {root / 'data'}") |
| return files |
|
|
|
|
| def file_sha256(path): |
| path = Path(path) |
| before = path.stat() |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for block in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(block) |
| after = path.stat() |
| if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns): |
| raise RuntimeError(f"File changed while hashing: {path}") |
| return digest.hexdigest() |
|
|
|
|
| def atomic_save(value, path): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(prefix=".pending-", dir=path.parent) as tmp: |
| pending = Path(tmp) / path.name |
| torch.save(value, pending) |
| os.replace(pending, path) |
|
|
|
|
| def _metadata(data_dir, split, n, seed, seq_len): |
| return { |
| "version": 1, "split": split, "seed": int(seed), |
| "seq_len": int(seq_len), "segments": int(n), |
| "encoding": "utf8-ignore+LF;byte-offset=2;overlap=1", |
| "files": [{"name": p.name, "sha256": file_sha256(p)} for p in parquet_files(data_dir, split)], |
| } |
|
|
|
|
| def _tensor_sha256(t): |
| return hashlib.sha256(t.numpy().astype("<i2", copy=False).tobytes()).hexdigest() |
|
|
|
|
| def _load_cache(path, metadata): |
| value = torch.load(path, map_location="cpu", weights_only=True) |
| if not isinstance(value, dict) or value.get("metadata") != metadata: |
| raise ValueError(f"Cache provenance is missing or mismatched: {path}. Use a new output directory or move the old cache aside.") |
| t = value.get("t") |
| if ( |
| not isinstance(t, torch.Tensor) or t.dtype != torch.int16 |
| or tuple(t.shape) != (metadata["segments"], metadata["seq_len"] + 1) |
| or int(t.min()) < 2 or int(t.max()) > 257 |
| or value.get("tensor_sha256") != _tensor_sha256(t) |
| ): |
| raise ValueError(f"Cache tensor failed integrity validation: {path}") |
| return t |
|
|
|
|
| def _segments(data_dir, split, seed, seq_len): |
| files = parquet_files(data_dir, split) |
| rng = random.Random(seed) |
| buf = bytearray() |
| while True: |
| order = list(files) |
| rng.shuffle(order) |
| nonempty = False |
| for path in order: |
| for rb in pq.ParquetFile(path).iter_batches(batch_size=256, columns=["text"]): |
| for text in rb.column(0).to_pylist(): |
| if not text: |
| continue |
| nonempty = True |
| buf.extend(text.encode("utf-8", errors="ignore")) |
| buf.extend(b"\n") |
| while len(buf) >= seq_len + 1: |
| yield torch.frombuffer(bytearray(buf[:seq_len + 1]), dtype=torch.uint8).long().add_(2) |
| del buf[:seq_len] |
| if not nonempty: |
| raise ValueError(f"No nonempty text in the {split} split.") |
|
|
|
|
| @torch.no_grad() |
| def build_eval_cache(data_dir, cache_dir, split, n, seed, seq_len, rank=0, world_size=1, distributed=False): |
| if n <= 0 or seq_len <= 0: |
| raise ValueError("Cache size and sequence length must be positive.") |
| root = Path(cache_dir) |
| path = root / f"slimpj_{split}_L{seq_len}_N{n}.pt" |
| status, tensor = [None], None |
| if not distributed or rank == 0: |
| try: |
| root.mkdir(parents=True, exist_ok=True) |
| with FileLock(str(path) + ".lock", timeout=300): |
| metadata = _metadata(data_dir, split, n, seed, seq_len) |
| if path.exists(): |
| tensor = _load_cache(path, metadata) |
| else: |
| stream = iter(_segments(data_dir, split, seed, seq_len)) |
| tensor = torch.stack([next(stream).to(torch.int16) for _ in range(n)]) |
| if _metadata(data_dir, split, n, seed, seq_len) != metadata: |
| raise RuntimeError("Dataset changed during cache generation.") |
| atomic_save({"t": tensor, "metadata": metadata, "tensor_sha256": _tensor_sha256(tensor)}, path) |
| status[0] = {"metadata": metadata} |
| except Exception as exc: |
| if not distributed: |
| raise |
| status[0] = {"error": f"{type(exc).__name__}: {exc}"} |
| if distributed: |
| dist.broadcast_object_list(status, src=0) |
| if "error" in status[0]: |
| raise RuntimeError(status[0]["error"]) |
| error = None |
| if rank != 0: |
| try: |
| tensor = _load_cache(path, status[0]["metadata"]) |
| except Exception as exc: |
| error = f"Rank {rank}: {type(exc).__name__}: {exc}" |
| errors = [None] * world_size |
| dist.all_gather_object(errors, error) |
| if any(errors): |
| raise RuntimeError("; ".join(x for x in errors if x)) |
| return tensor |
|
|
|
|
| def hankel_kernel(length, eps=1e-6): |
| s = torch.arange(2, 2 * length + 1, dtype=torch.float32, device="cpu") |
| a = 2 / (s ** 3 - s + eps) |
| return torch.cat([a[length - 1:], torch.zeros(1), a[:length - 1]]) |
|
|
|
|
| def hmv(g_f, x): |
| length, p = x.shape |
| padded = torch.zeros(p, 2 * length, dtype=x.dtype, device=x.device) |
| padded[:, :length] = x.flip(0).T |
| return torch.fft.irfft(torch.fft.rfft(padded, dim=-1) * g_f[None], n=2 * length, dim=-1)[:, :length].T |
|
|
|
|
| def _basis_digest(sigma, phi): |
| digest = hashlib.sha256() |
| for tensor in (sigma, phi): |
| digest.update(tensor.numpy().astype("<f4", copy=False).tobytes()) |
| return digest.hexdigest() |
|
|
|
|
| @torch.no_grad() |
| def hankel_basis(cache_dir, length, channels, iters=60, oversample=16, tol=1e-6): |
| if not 1 <= channels <= length or iters <= 0 or oversample < 0 or tol <= 0: |
| raise ValueError("Invalid Hankel basis configuration.") |
| metadata = dict(version=2, length=length, channels=channels, iters=iters, oversample=oversample, tol=tol, seed=0) |
| signature = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode()).hexdigest()[:12] |
| root = Path(cache_dir) |
| root.mkdir(parents=True, exist_ok=True) |
| path = root / f"hankel_v2_L{length}_K{channels}_{signature}.pt" |
| with FileLock(str(path) + ".lock", timeout=300): |
| if path.exists(): |
| saved = torch.load(path, map_location="cpu", weights_only=True) |
| sigma, phi = saved["sigma"], saved["phi"] |
| if ( |
| saved.get("metadata") != metadata or sigma.shape != (channels,) |
| or phi.shape != (length, channels) |
| or not torch.isfinite(sigma).all() or not torch.isfinite(phi).all() |
| or saved.get("sha256") != _basis_digest(sigma, phi) |
| ): |
| raise ValueError(f"Invalid Hankel cache: {path}") |
| return sigma, phi |
| generator = torch.Generator(device="cpu").manual_seed(0) |
| q = torch.randn(length, channels + oversample, generator=generator, device="cpu", dtype=torch.float32) |
| q, _ = torch.linalg.qr(q, mode="reduced") |
| gf = torch.fft.rfft(hankel_kernel(length)) |
| for step in range(1, iters + 1): |
| q, _ = torch.linalg.qr(hmv(gf, q), mode="reduced") |
| if step % max(1, iters // 4) == 0 or step == iters: |
| values, vectors = torch.linalg.eigh(q.T @ hmv(gf, q)) |
| order = values.argsort(descending=True) |
| values, vectors = values[order], vectors[:, order] |
| q = q @ vectors |
| residual = (hmv(gf, q[:, :channels]) - q[:, :channels] * values[:channels][None]).norm(dim=0) |
| if float((residual / (values[:channels].abs() + 1e-12)).max()) < tol: |
| break |
| values, vectors = torch.linalg.eigh(q.T @ hmv(gf, q)) |
| order = values.argsort(descending=True) |
| sigma, phi = values[order][:channels], (q @ vectors[:, order])[:, :channels] |
| atomic_save({"sigma": sigma, "phi": phi, "metadata": metadata, "sha256": _basis_digest(sigma, phi)}, path) |
| return sigma, phi |
|
|
