| """Corpus download (pinned + licensed), text pretokenization, dataset validation and statistics.
|
|
|
| Every external corpus is listed in SOURCES with a pinned revision and its license, and the
|
| provenance record is written next to the downloaded files so caches can be audited later.
|
| """
|
| from __future__ import annotations
|
|
|
| import hashlib
|
| import json
|
| from pathlib import Path
|
|
|
| import numpy as np
|
| from PIL import Image
|
|
|
| SOURCES: dict[str, dict] = {
|
| "tinystories": {
|
| "repo_id": "roneneldan/TinyStories",
|
| "repo_type": "dataset",
|
| "revision": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
|
| "files": [
|
| "data/train-00000-of-00004-2d5a1467fff1081b.parquet",
|
| "data/validation-00000-of-00001-869c898b519ad725.parquet",
|
| ],
|
| "license": "CDLA-Sharing-1.0",
|
| "url": "https://huggingface.co/datasets/roneneldan/TinyStories",
|
| "notes": "Synthetic children's stories (GPT-3.5/4 generated). Only 1 of 4 train shards is used.",
|
| },
|
| "dolly": {
|
| "repo_id": "databricks/databricks-dolly-15k",
|
| "repo_type": "dataset",
|
| "revision": "bdd27f4d94b9c1f951818a7da7fd7aeea5dbff1a",
|
| "files": ["databricks-dolly-15k.jsonl"],
|
| "license": "CC-BY-SA-3.0",
|
| "url": "https://huggingface.co/datasets/databricks/databricks-dolly-15k",
|
| "notes": "Human-written instruction/response pairs; used for instruction tuning only.",
|
| },
|
| }
|
|
|
|
|
| def sha256_file(path: str | Path) -> str:
|
| h = hashlib.sha256()
|
| with open(path, "rb") as f:
|
| for chunk in iter(lambda: f.read(1 << 20), b""):
|
| h.update(chunk)
|
| return h.hexdigest()
|
|
|
|
|
| def download_sources(raw_dir: str | Path, names: list[str] | None = None) -> dict[str, list[Path]]:
|
| """Download the pinned corpus files and write a provenance.json record."""
|
| from huggingface_hub import hf_hub_download
|
|
|
| raw_dir = Path(raw_dir)
|
| out: dict[str, list[Path]] = {}
|
| for name in names or list(SOURCES):
|
| src = SOURCES[name]
|
| dest = raw_dir / name
|
| paths = [
|
| Path(hf_hub_download(src["repo_id"], f, repo_type=src["repo_type"], revision=src["revision"], local_dir=dest))
|
| for f in src["files"]
|
| ]
|
| record = {k: v for k, v in src.items()} | {"sha256": {p.name: sha256_file(p) for p in paths}}
|
| (dest / "provenance.json").write_text(json.dumps(record, indent=2))
|
| out[name] = paths
|
| return out
|
|
|
|
|
| def load_tinystories(raw_dir: str | Path, split: str) -> list[str]:
|
| import pyarrow.parquet as pq
|
|
|
| files = sorted((Path(raw_dir) / "tinystories" / "data").glob(f"{split}-*.parquet"))
|
| if not files:
|
| raise FileNotFoundError(f"No TinyStories {split} parquet under {raw_dir}; run scripts/generate_dataset.py download")
|
| texts: list[str] = []
|
| for f in files:
|
| texts += [t.strip() for t in pq.read_table(f, columns=["text"]).column("text").to_pylist() if t and t.strip()]
|
| return texts
|
|
|
|
|
| def load_dolly(raw_dir: str | Path) -> list[dict]:
|
| rows = []
|
| with open(Path(raw_dir) / "dolly" / "databricks-dolly-15k.jsonl", encoding="utf-8") as f:
|
| for line in f:
|
| r = json.loads(line)
|
| if r["instruction"].strip() and r["response"].strip():
|
| rows.append(r)
|
| return rows
|
|
|
|
|
| def tokenize_stream(texts: list[str], encode_batch, prefix: list[int], eos_id: int, out_path: str | Path) -> int:
|
| """Concatenate documents as <prefix> doc [EOS] ... into one uint16 token stream (for packing)."""
|
| out_path = Path(out_path)
|
| out_path.parent.mkdir(parents=True, exist_ok=True)
|
| n = 0
|
| with open(out_path, "wb") as f:
|
| for i in range(0, len(texts), 4096):
|
| arrs = [np.array([*prefix, *ids, eos_id], dtype=np.uint16) for ids in encode_batch(texts[i : i + 4096])]
|
| buf = np.concatenate(arrs)
|
| buf.tofile(f)
|
| n += len(buf)
|
| return n
|
|
|
|
|
| def validate_sequence(tokens: np.ndarray, vocab) -> None:
|
| """Raise ValueError if a token sequence violates the unified-vocabulary structure rules."""
|
| vocab.check_structure(tokens)
|
|
|
|
|
| def dataset_stats(cache_dir: str | Path) -> dict:
|
| """Summarise every cache file (token counts / sample counts) for docs and the dashboard."""
|
| cache_dir = Path(cache_dir)
|
| stats: dict[str, dict] = {}
|
| for p in sorted(cache_dir.glob("*.bin")):
|
| stats[p.name] = {"tokens": p.stat().st_size // 2}
|
| for p in sorted(cache_dir.glob("*.npz")):
|
| with np.load(p) as z:
|
| stats[p.name] = {k: list(z[k].shape) for k in z.files}
|
| return stats
|
|
|
|
|
|
|
| def prepare_image(img, res: int = 64) -> np.ndarray:
|
| """PIL / array / path -> uint8 [res,res,3] (centre square crop, area resize)."""
|
| if isinstance(img, (str, Path)):
|
| img = Image.open(img)
|
| if isinstance(img, np.ndarray):
|
| img = Image.fromarray(img.astype(np.uint8))
|
| img = img.convert("RGB")
|
| w, h = img.size
|
| s = min(w, h)
|
| img = img.crop(((w - s) // 2, (h - s) // 2, (w - s) // 2 + s, (h - s) // 2 + s))
|
| return np.asarray(img.resize((res, res), Image.BOX if s >= res else Image.BICUBIC))
|
|
|
|
|
| def load_video(src, max_frames: int = 8, fps: float = 4.0, res: int = 64) -> np.ndarray:
|
| """Video file (mp4/gif/webm...) or frame list -> uint8 [T,res,res,3], resampled to ~`fps`."""
|
| if not isinstance(src, (str, Path)):
|
| return np.stack([prepare_image(f, res) for f in list(src)[:max_frames]])
|
| import imageio_ffmpeg
|
|
|
| reader = imageio_ffmpeg.read_frames(str(src))
|
| meta = next(reader)
|
| w, h = meta["size"]
|
| stride = max(1, round(float(meta.get("fps") or fps) / fps))
|
| frames = []
|
| for i, raw in enumerate(reader):
|
| if i % stride == 0:
|
| frames.append(prepare_image(np.frombuffer(raw, np.uint8).reshape(h, w, 3), res))
|
| if len(frames) == max_frames:
|
| break
|
| reader.close()
|
| if not frames:
|
| raise ValueError(f"could not decode any frames from {src}")
|
| return np.stack(frames)
|
|
|
|
|
| def video_windows(path: str | Path, n_frames: int = 8, fps: float = 4.0, max_windows: int = 4, res: int = 64) -> list[np.ndarray]:
|
| """Cut a video file into up to `max_windows` consecutive `n_frames`-frame clips at ~`fps`.
|
| Clips too short for one window are sampled evenly at native frame rate; < n_frames frames -> []."""
|
| clip = load_video(path, max_frames=n_frames * max_windows, fps=fps, res=res)
|
| if len(clip) >= n_frames:
|
| return [clip[i : i + n_frames] for i in range(0, len(clip) - n_frames + 1, n_frames)]
|
| native = load_video(path, max_frames=10_000, fps=1e9, res=res)
|
| if len(native) < n_frames:
|
| return []
|
| return [native[np.linspace(0, len(native) - 1, n_frames).round().astype(int)]]
|
|
|
|
|
|
|
| IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
| VIDEO_EXT = {".mp4", ".mov", ".webm", ".mkv", ".avi", ".gif"}
|
| USER_KEYS = ("user", "instruction", "prompt", "question", "input")
|
| ASSISTANT_KEYS = ("assistant", "response", "output", "answer", "completion")
|
|
|
|
|
| def scan_user_dataset(root: str | Path) -> dict[str, list]:
|
| """Collect samples from a user-provided folder or .jsonl file, routing records by their keys.
|
|
|
| JSONL records (paths relative to the .jsonl file):
|
| {"text": ...} -> text corpus
|
| {"image": p, "caption": ...} -> image-caption (t2i + i2t)
|
| {"image": p, "target_image": p, "instruction": ...} -> image editing (i2i)
|
| {"video": p, "caption": ...} (caption optional) -> video-caption (t2v, v2t) / continuation
|
| {"instruction"|"prompt"|"user": ..., "response"|"output"|"assistant": ...} -> instruction tuning
|
| {"messages": [{"role": "user", ...}, {"role": "assistant", ...}, ...]} -> instruction tuning
|
| Folders: *.txt -> text documents; image/video files with a same-name .txt sidecar -> captioned
|
| media (images without a caption are skipped; videos without one are kept for continuation).
|
| """
|
| root = Path(root)
|
| out: dict[str, list] = {"texts": [], "images": [], "edits": [], "videos": [], "chats": [], "skipped": []}
|
| files = [root] if root.is_file() else sorted(p for p in root.rglob("*") if p.is_file())
|
| sidecar = {p.with_suffix(".txt") for p in files if p.suffix.lower() in IMAGE_EXT | VIDEO_EXT}
|
| for p in files:
|
| if p.suffix.lower() == ".jsonl":
|
| for ln, line in enumerate(p.read_text(encoding="utf-8", errors="replace").splitlines()):
|
| if not line.strip():
|
| continue
|
| try:
|
| _route(json.loads(line), p.parent, out)
|
| except (ValueError, KeyError, TypeError) as e:
|
| out["skipped"].append(f"{p.name}:{ln + 1}: {e}")
|
|
|
| referenced = {Path(x).resolve() for k in ("images", "videos") for x, _ in out[k]} | {Path(x).resolve() for e in out["edits"] for x in e[:2]}
|
| for p in files:
|
| ext = p.suffix.lower()
|
| if ext == ".jsonl" or p.resolve() in referenced:
|
| continue
|
| if ext == ".txt" and p not in sidecar:
|
| t = p.read_text(encoding="utf-8", errors="replace").strip()
|
| if t:
|
| out["texts"].append(t)
|
| elif ext in IMAGE_EXT | VIDEO_EXT:
|
| cap_p = p.with_suffix(".txt")
|
| cap = cap_p.read_text(encoding="utf-8", errors="replace").strip() if cap_p.exists() else ""
|
| if ext in VIDEO_EXT:
|
| out["videos"].append((p, cap))
|
| elif cap:
|
| out["images"].append((p, cap))
|
| else:
|
| out["skipped"].append(f"{p.name}: image without caption")
|
| return out
|
|
|
|
|
| def _route(r: dict, base: Path, out: dict) -> None:
|
| get = lambda keys: next((str(r[k]).strip() for k in keys if r.get(k)), "")
|
| if "messages" in r:
|
| msgs = r["messages"]
|
| for a, b in zip(msgs, msgs[1:]):
|
| if a.get("role") == "user" and b.get("role") == "assistant" and a.get("content") and b.get("content"):
|
| out["chats"].append((str(a["content"]).strip(), str(b["content"]).strip()))
|
| elif "video" in r:
|
| out["videos"].append((base / r["video"], get(("caption", "text", "prompt"))))
|
| elif "image" in r and ("target_image" in r or "edited_image" in r):
|
| out["edits"].append((base / r["image"], base / (r.get("target_image") or r["edited_image"]), get(("instruction", "edit", "prompt"))))
|
| elif "image" in r:
|
| cap = get(("caption", "text", "prompt"))
|
| if not cap:
|
| raise ValueError("image record without caption")
|
| out["images"].append((base / r["image"], cap))
|
| elif get(USER_KEYS) and get(ASSISTANT_KEYS):
|
| user = get(USER_KEYS) + (("\n\n" + str(r["context"]).strip()) if r.get("context") else "")
|
| out["chats"].append((user, get(ASSISTANT_KEYS)))
|
| elif get(("text", "content", "story")):
|
| out["texts"].append(get(("text", "content", "story")))
|
| else:
|
| raise ValueError(f"unrecognised record keys {sorted(r)}")
|
|
|
|
|
| def split_of(key: str, val_pct: float) -> str:
|
| """Deterministic train/val assignment from a content hash (stable across runs and machines)."""
|
| return "val" if int(hashlib.sha1(key.encode("utf-8", "replace")).hexdigest()[:8], 16) % 10_000 < val_pct * 100 else "train"
|
|
|
