v18 — natural-language captions, think-label sidecars, parquet support

Three changes that move the trainer from "hard-coded style trigger
prepended at training time" to "style anchor lives inside the natural
language caption + matched think-mode reasoning":

1. **Caption + think in a single .txt** (compact format):
caption_paragraph
---think---
6-section reasoning_paragraph
Parsed by `parse_caption_and_think()`. Caption-only files (no marker)
stay valid. Legacy split-format with `<id>.think.txt` is still read
when no marker is found, so existing v18 datasets work without
migration.

2. **`<think>` injection in the prompt template**: when `T2ISample.think`
is non-empty, `SenseNovaU1Collator` renders it inside the otherwise-
empty `<think></think>` block of the official chat template. This
matches the inference distribution when sampling with `--think-mode`,
where the model autoregressively fills the same window — fixes the
"noise output at sample time after train didn't see filled-think"
class of failure.

3. **`ArrowT2IDataset`** for 1M-scale runs: parquet shards with schema
`(sample_id, caption, think?, image bytes or path)`. Trainer
auto-dispatches based on `data.data_dir` suffix. New CLI:
python -m train_u1.scripts.dataset_tools pack-arrow FOLDER --out OUT.parquet
python -m train_u1.scripts.dataset_tools inspect-arrow OUT.parquet
python -m train_u1.scripts.dataset_tools unify-txt FOLDER # fold .txt + .think.txt → unified .txt

`configs/v18.yaml` ships the new recipe (empty trigger, conservative
VRAM settings for the longer prefix from think labels).

Adds `pyarrow>=14` to deps. 10 new tests in `test_dataset_format.py`
cover the parser, both legacy + unified PairedFolderT2IDataset paths,
and an arrow round-trip.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

README.md

@@ -74,7 +74,9 @@ HF_HOME=$PWD/hf_cache python -m train_u1.scripts.install_modeling_into_snapshot
 
 ## Train
 
-1. Lay out your data:
+1. Lay out your data. Two formats are supported.
+
+   **Folder of paired files** (recommended for ≤ ~10k images):
 
    ```
    dataset/my_style/
@@ -83,9 +85,44 @@ HF_HOME=$PWD/hf_cache python -m train_u1.scripts.install_modeling_into_snapshot
    └── …           └── …
    ```
 
-   Each `.txt` is a plain-text caption. The style trigger string (`my style`,
-   `hayateluc style`, …) is **prepended** at training time — don't bake it
-   into the captions yourself.
+   Each `.txt` is a single-paragraph natural-language caption. Embed the
+   artist credit / style anchor inside the description naturally — don't
+   rely on a hard-coded trigger prepend (`style.trigger` in the YAML is
+   for backward compat only; the v18 recipe uses an empty trigger).
+
+   **Optional**: append a `<think>...</think>` reasoning label inside the
+   same `.txt` after a `---think---` delimiter line:
+
+   ```
+   An illustration by Hayateluc depicting a wisteria-trellis path under
+   morning glow, painterly composition, no people.
+   ---think---
+   1. **Instruction Understanding:** ...
+   6. **Explicit Prompt:** ...
+   ```
+
+   When present, the trainer renders this into the prompt template's
+   `<think>` window so train-time distribution matches inference
+   `--think-mode` (avoids prefix-distribution shift on long autoregressive
+   think). For batch generation of think labels see Agent B's prompt in
+   the v18 commit history; or write them yourself in the upstream
+   6-section format.
+
+   **Parquet/arrow shards** (recommended for ≥ ~10k images, e.g. 1M
+   scaling):
+
+   ```bash
+   # Pack a folder dataset → single parquet shard
+   python -m train_u1.scripts.dataset_tools pack-arrow dataset/my_style \
+       --out artifacts/my_style.parquet
+   # Inspect first 3 rows
+   python -m train_u1.scripts.dataset_tools inspect-arrow artifacts/my_style.parquet
+   ```
+
+   Schema: `sample_id, caption, think (nullable), image (binary)`. Set
+   `data.data_dir` in the YAML to point at the parquet path; the training
+   script auto-detects `.parquet` and uses `ArrowT2IDataset` instead of
+   `PairedFolderT2IDataset`.
 
 2. Edit `configs/default.yaml`. The only fields you must touch:
 

configs/v18.yaml

@@ -0,0 +1,56 @@
+# v18 — natural-language captions + per-sample think labels.
+#
+# Drops the hard-coded `style.trigger: "hayateluc style, "` prepend that v16c
+# was using, in favour of:
+#   - rewritten captions in `dataset/Hayateluc/<id>.txt` that embed the
+#     "Hayateluc / in the style of Hayateluc" anchor naturally inside the
+#     description (varied across the 56 samples)
+#   - per-sample `<id>.think.txt` sidecars containing a 6-section reasoning
+#     block whose section 5 explicitly anchors "Hayateluc's painterly anime
+#     style". The collator injects this into the prompt template's empty
+#     `<think>...</think>` window so train- and infer-with-think distributions
+#     match.
+#
+# Starts from the same trainable surface as v16c (LoRA attn+mlp r=64 +
+# ts/ns/vision/fm_head full-FT). Re-trained from scratch since the prefix
+# distribution changed — old v16c LoRA isn't a useful warm-start.
+
+run_name: v18
+
+data:
+  data_dir: dataset/Hayateluc
+  cap_max_pixels: 4194304
+  snap_bucket: true
+
+style:
+  trigger: ""                        # ← drop the hard-coded prepend
+  prompt_template: official
+
+lora:
+  spec: "attn=r64a64;mlp=r64a64"
+  dropout: 0.0
+
+unfreeze:
+  - '^fm_modules\.timestep_embedder\.'
+  - '^fm_modules\.noise_scale_embedder\.'
+  - '^fm_modules\.vision_model_mot_gen\.'
+  - '^fm_modules\.fm_head\.'
+
+train:
+  steps: 6000
+  lr: 5.0e-5
+  seed: 0
+  shuffle: true
+  grad_accum: 1
+  checkpoint_every: 600
+
+runtime:
+  # think labels add ~320 tokens per sample's prefix → 56 × 320 × 172KB ≈ +3 GB
+  # of prefix KV. With gc_skip_last=6's activation residency stack (which used
+  # to fit at v16c's shorter prefix) this OOMs at ~30 GB on a 32 GB card.
+  # Drop both: stream KVs from CPU per-step + full GC on all decoder layers.
+  # Tradeoff is ~30% slower per step, but stable.
+  keep_kvs_on_gpu: false
+  gc_skip_last: 0
+  device: cuda
+  cpu_device: cpu

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
   "peft>=0.13",
   "pyyaml",
   "pydantic>=2.0",
+  "pyarrow>=14",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -20,9 +20,10 @@ numpy
 bitsandbytes>=0.45    # 8-bit / 4-bit base + paged AdamW8bit
 peft>=0.13            # not strictly required (we have our own LoRA), pulled in for utility helpers
 
-# config
+# config + dataset
 pyyaml
 pydantic>=2.0
+pyarrow>=14         # ArrowT2IDataset (parquet shards) for 1M+ scale datasets
 
 # tests
 pytest>=7

train_u1/data/collators.py

@@ -204,14 +204,28 @@ class SenseNovaU1Collator:
         else:
             raw_prompts = [s.prompt for s in samples]
         if self._build_t2i_query is not None:
-            prompts = [
-                self._build_t2i_query(
-                    rp,
-                    system_message=self._sys_msg_for_gen,
-                    append_text=self._gen_append,
+            prompts = []
+            for rp, s in zip(raw_prompts, samples):
+                # Per-sample think injection: when the dataset supplies a
+                # `think` text, render it INSIDE the otherwise-empty
+                # `<think></think>` block of the official prompt template.
+                # This makes training distribution match inference-time
+                # `--think-mode`, where the model autoregressively fills the
+                # same window with ~250-400 reasoning tokens. Without this,
+                # the gen tower sees an unfamiliar prefix length/content
+                # at inference and the LoRA delta is calibrated against
+                # the wrong cond-KV distribution.
+                if s.think:
+                    append_text = f"<think>\n{s.think}\n</think>\n\n<img>"
+                else:
+                    append_text = self._gen_append
+                prompts.append(
+                    self._build_t2i_query(
+                        rp,
+                        system_message=self._sys_msg_for_gen,
+                        append_text=append_text,
+                    )
                 )
-                for rp in raw_prompts
-            ]
         else:
             prompts = list(raw_prompts)
         input_ids, text_lens = self._tokenize(prompts)

train_u1/data/datasets.py

@@ -32,6 +32,12 @@ class T2ISample:
     prompt: str
     image: torch.Tensor      # (3, H, W) in [0, 1] or normalized — collator decides
     seed: int = 0
+    # Optional pre-computed `<think>...</think>` reasoning text. When set, the
+    # collator embeds it INSIDE the empty think block of the official prompt
+    # template, so training distribution matches inference-time `--think-mode`.
+    # When None, the empty `<think>\n\n</think>` block is preserved (matches
+    # inference-time without `--think-mode`).
+    think: str | None = None
 
 
 class SyntheticT2ITinyDataset(Dataset):
@@ -164,13 +170,168 @@ class PairedFolderT2IDataset(Dataset):
 
         img_path, txt_path, stem = self.pairs[idx]
         with open(txt_path, encoding="utf-8") as f:
-            caption = f.read().strip()
+            raw = f.read()
+        caption, think_text = parse_caption_and_think(raw)
         if self.prompt_template:
             caption = self.prompt_template.format(caption=caption)
+        # Legacy fallback: `<id>.think.txt` separate sidecar (deprecated;
+        # `parse_caption_and_think` is the preferred path).
+        if think_text is None:
+            think_path = img_path.with_suffix(".think.txt")
+            if think_path.is_file():
+                with open(think_path, encoding="utf-8") as f:
+                    think_text = f.read().strip() or None
         chw, _hw = load_and_preprocess_image(
             img_path,
             cap_max_pixels=self.cap_max_pixels,
             normalize="x0",  # 公开证据显示 — fm_head output space, [-1, 1] (NOT ImageNet)
             snap_bucket=self.snap_bucket,
         )
-        return T2ISample(sample_id=stem, prompt=caption, image=chw, seed=hash(stem) & 0xFFFF)
+        return T2ISample(
+            sample_id=stem, prompt=caption, image=chw,
+            seed=hash(stem) & 0xFFFF, think=think_text,
+        )
+
+
+# Marker used inside a single `.txt` to separate caption from think label.
+# Any whitespace tolerated around the marker; case-insensitive on the
+# `THINK` keyword. Pick a marker unlikely to appear in natural prose.
+THINK_DELIMITER_RE = __import__("re").compile(r"^\s*-{3,}\s*think\s*-{3,}\s*$", flags=__import__("re").IGNORECASE | __import__("re").MULTILINE)
+
+
+def parse_caption_and_think(raw: str) -> tuple[str, str | None]:
+    """Split a caption file into `(caption, think)` parts.
+
+    The new compact format puts both labels in a single `.txt`::
+
+        A natural-language caption embedding the artist style.
+        ---think---
+        1. **Instruction Understanding:** ...
+        ...
+        6. **Explicit Prompt:** ...
+
+    Falls back to "whole file is caption, no think" for plain captions.
+    Returns `(caption, None)` if no marker is present, else `(caption, think)`
+    with both stripped.
+    """
+    m = THINK_DELIMITER_RE.search(raw)
+    if m is None:
+        return raw.strip(), None
+    caption = raw[: m.start()].strip()
+    think = raw[m.end():].strip()
+    return caption, (think or None)
+
+
+# --------------------------------------------------------------------------- #
+# ArrowT2IDataset (large-scale)                                               #
+# --------------------------------------------------------------------------- #
+
+
+class ArrowT2IDataset(Dataset):
+    """T2I dataset backed by a parquet/arrow shard.
+
+    Schema expected (rows aligned with `T2ISample` fields):
+        - `sample_id` : string
+        - `caption`   : string
+        - `think`     : string (optional; nullable column)
+        - `image`     : binary (raw image bytes, e.g. PNG/JPEG) — preferred
+                         OR `image_path` : string (path resolved relative to
+                         `image_root` if relative, else absolute)
+
+    Reads via `pyarrow` table-of-shards mmap, so memory-efficient even at
+    millions of rows. Suitable for the 1M-image scaling experiment (task #53).
+
+    Args:
+        path: parquet file or directory of parquet shards
+        image_root: base dir to resolve relative `image_path` columns
+        cap_max_pixels: optional VRAM-friendly cap (passed to smart_resize)
+        prompt_template: optional template (matches PairedFolderT2IDataset)
+        snap_bucket: snap to nearest official bucket
+    """
+
+    def __init__(
+        self,
+        path: str | os.PathLike,
+        *,
+        image_root: str | os.PathLike | None = None,
+        cap_max_pixels: int | None = None,
+        prompt_template: str | None = None,
+        snap_bucket: bool = False,
+    ):
+        try:
+            import pyarrow.parquet as pq  # noqa: F401
+        except ImportError as e:
+            raise ImportError(
+                "ArrowT2IDataset requires pyarrow. Install with `pip install pyarrow>=14`."
+            ) from e
+        self.path = Path(path)
+        self.image_root = Path(image_root) if image_root else None
+        self.cap_max_pixels = cap_max_pixels
+        self.prompt_template = prompt_template
+        self.snap_bucket = snap_bucket
+        self._table = None  # lazy-loaded
+        self._n: int | None = None
+
+    def _ensure_table(self):
+        if self._table is not None:
+            return
+        import pyarrow.parquet as pq
+        if self.path.is_dir():
+            # Directory of shards — read combined.
+            self._table = pq.read_table(str(self.path))
+        else:
+            self._table = pq.read_table(str(self.path))
+        cols = self._table.column_names
+        # Schema sanity
+        if "sample_id" not in cols or "caption" not in cols:
+            raise RuntimeError(
+                f"{self.path}: required columns missing. "
+                f"Schema must include `sample_id` and `caption`. Got: {cols}"
+            )
+        if "image" not in cols and "image_path" not in cols:
+            raise RuntimeError(
+                f"{self.path}: must have either `image` (bytes) or `image_path` (string). Got: {cols}"
+            )
+        self._n = self._table.num_rows
+
+    def __len__(self) -> int:
+        self._ensure_table()
+        return self._n  # type: ignore[return-value]
+
+    def __getitem__(self, idx: int) -> T2ISample:
+        from io import BytesIO
+
+        from train_u1.data.u1_preprocess import (
+            load_and_preprocess_image,
+            preprocess_pil_image,
+        )
+
+        self._ensure_table()
+        row = self._table.slice(idx, 1).to_pydict()
+        sample_id = row["sample_id"][0]
+        caption = row["caption"][0]
+        think = (row.get("think") or [None])[0] or None
+
+        if "image" in self._table.column_names and row["image"][0] is not None:
+            from PIL import Image
+            pil = Image.open(BytesIO(row["image"][0])).convert("RGB")
+            chw, _hw = preprocess_pil_image(
+                pil, cap_max_pixels=self.cap_max_pixels,
+                normalize="x0", snap_bucket=self.snap_bucket,
+            )
+        else:
+            img_path_str = row["image_path"][0]
+            img_path = Path(img_path_str)
+            if not img_path.is_absolute() and self.image_root is not None:
+                img_path = self.image_root / img_path
+            chw, _hw = load_and_preprocess_image(
+                img_path, cap_max_pixels=self.cap_max_pixels,
+                normalize="x0", snap_bucket=self.snap_bucket,
+            )
+
+        if self.prompt_template:
+            caption = self.prompt_template.format(caption=caption)
+        return T2ISample(
+            sample_id=str(sample_id), prompt=str(caption), image=chw,
+            seed=hash(sample_id) & 0xFFFF, think=think,
+        )

train_u1/data/u1_preprocess.py

@@ -101,8 +101,8 @@ def snap_to_official_bucket(H: int, W: int) -> tuple[int, int]:
     return best
 
 
-def load_and_preprocess_image(
-    path: str | Path,
+def preprocess_pil_image(
+    img: Image.Image,
     *,
     factor: int = PATCH32,
     min_pixels: int = SMART_RESIZE_MIN_PIXELS,
@@ -111,26 +111,12 @@ def load_and_preprocess_image(
     normalize: str = "x0",
     snap_bucket: bool = False,
 ) -> tuple[torch.Tensor, tuple[int, int]]:
-    """Read RGB image, smart-resize, return `(image_chw, (H, W))`.
-
-    `normalize`:
-    - `"x0"`     → mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5) → output in [-1, 1].
-                   Use this for FM training x0 targets (matches fm_head output
-                   space + `t2i_generate.image_prediction` Euler state).
-    - `"vision"` → ImageNet mean/std. Use this for ordinary `vision_model`
-                   inputs (understanding path).
-    - `"none"`   → raw [0, 1] (debugging only).
-
-    `cap_max_pixels` further clamps the upstream max (e.g. 4194304 for the
-    2048² training bucket).
+    """In-memory variant of `load_and_preprocess_image` taking a PIL.Image.
 
-    `snap_bucket`: if True, pick the closest aspect-ratio match in
-    `OFFICIAL_BUCKETS_HW` (overrides smart_resize/min/max/cap). Use this for
-    training when you want every sample to land on a known-supported bucket
-    shape, eliminating the "trained on arbitrary shape, sampled at 2048²"
-    mismatch.
+    Use this for arrow/parquet datasets that decode image bytes inline.
+    Returns `(image_chw, (H, W))` with the same conventions.
     """
-    img = Image.open(path).convert("RGB")
+    img = img.convert("RGB")
     W, H = img.size
     if snap_bucket:
         H_bar, W_bar = snap_to_official_bucket(H, W)
@@ -141,11 +127,11 @@ def load_and_preprocess_image(
 
     import numpy as np
 
-    arr = np.asarray(img).astype("float32") / 255.0  # (H, W, 3) in [0, 1]
+    arr = np.asarray(img).astype("float32") / 255.0
     if normalize == "x0":
         mean = np.array(X0_MEAN, dtype="float32")
         std = np.array(X0_STD, dtype="float32")
-        arr = (arr - mean) / std  # → [-1, 1]
+        arr = (arr - mean) / std
     elif normalize == "vision":
         mean = np.array(IMAGENET_MEAN, dtype="float32")
         std = np.array(IMAGENET_STD, dtype="float32")
@@ -155,5 +141,41 @@ def load_and_preprocess_image(
     else:
         raise ValueError(f"normalize must be 'x0' / 'vision' / 'none', got {normalize!r}")
 
-    chw = torch.from_numpy(arr.transpose(2, 0, 1))  # (3, H, W)
+    chw = torch.from_numpy(arr.transpose(2, 0, 1))
     return chw, (H_bar, W_bar)
+
+
+def load_and_preprocess_image(
+    path: str | Path,
+    *,
+    factor: int = PATCH32,
+    min_pixels: int = SMART_RESIZE_MIN_PIXELS,
+    max_pixels: int = SMART_RESIZE_MAX_PIXELS,
+    cap_max_pixels: int | None = None,
+    normalize: str = "x0",
+    snap_bucket: bool = False,
+) -> tuple[torch.Tensor, tuple[int, int]]:
+    """Read RGB image, smart-resize, return `(image_chw, (H, W))`.
+
+    `normalize`:
+    - `"x0"`     → mean=(0.5,0.5,0.5), std=(0.5,0.5,0.5) → output in [-1, 1].
+                   Use this for FM training x0 targets (matches fm_head output
+                   space + `t2i_generate.image_prediction` Euler state).
+    - `"vision"` → ImageNet mean/std. Use this for ordinary `vision_model`
+                   inputs (understanding path).
+    - `"none"`   → raw [0, 1] (debugging only).
+
+    `cap_max_pixels` further clamps the upstream max (e.g. 4194304 for the
+    2048² training bucket).
+
+    `snap_bucket`: if True, pick the closest aspect-ratio match in
+    `OFFICIAL_BUCKETS_HW` (overrides smart_resize/min/max/cap). Use this for
+    training when you want every sample to land on a known-supported bucket
+    shape, eliminating the "trained on arbitrary shape, sampled at 2048²"
+    mismatch.
+    """
+    img = Image.open(path)
+    return preprocess_pil_image(
+        img, factor=factor, min_pixels=min_pixels, max_pixels=max_pixels,
+        cap_max_pixels=cap_max_pixels, normalize=normalize, snap_bucket=snap_bucket,
+    )

train_u1/scripts/dataset_tools.py

@@ -0,0 +1,202 @@
+"""Dataset format conversion utilities.
+
+Three subcommands:
+
+    unify-txt    — fold legacy `<id>.txt` + `<id>.think.txt` into a single
+                   `<id>.txt` with the `---think---` delimiter, in place.
+                   Use after a successful training run when you want to
+                   simplify the on-disk layout. Originals at `<id>.old.txt`
+                   are left untouched.
+
+    pack-arrow   — pack a folder dataset into a parquet shard for the
+                   ArrowT2IDataset path. Schema written:
+                       sample_id, caption, think (nullable), image (bytes)
+                   Suitable for 1M-image scaling — one shard per ~10-50k images.
+
+    inspect-arrow — print row count + first 3 rows of a parquet shard for
+                   sanity checking.
+
+Usage:
+    python -m train_u1.scripts.dataset_tools unify-txt dataset/Hayateluc
+    python -m train_u1.scripts.dataset_tools pack-arrow dataset/Hayateluc \
+        --out artifacts/hayateluc.parquet
+    python -m train_u1.scripts.dataset_tools inspect-arrow artifacts/hayateluc.parquet
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+
+THINK_MARKER = "---think---"
+
+
+def _read(p: Path) -> str:
+    with open(p, encoding="utf-8") as f:
+        return f.read().strip()
+
+
+def cmd_unify_txt(folder: Path, *, dry_run: bool = False) -> int:
+    """Fold `<id>.txt` + `<id>.think.txt` → unified `<id>.txt` (in place).
+
+    Skips ids whose `<id>.txt` already contains the `---think---` marker
+    (idempotent). Skips ids without a corresponding `<id>.think.txt`.
+    """
+    if not folder.is_dir():
+        print(f"folder not found: {folder}", file=sys.stderr); return 2
+
+    n_unified = 0
+    n_skipped_already = 0
+    n_skipped_no_think = 0
+    for txt in sorted(folder.glob("*.txt")):
+        if txt.name.endswith(".old.txt") or txt.name.endswith(".think.txt"):
+            continue
+        body = _read(txt)
+        if THINK_MARKER in body.lower():
+            n_skipped_already += 1
+            continue
+        think_path = txt.with_suffix(".think.txt")
+        if not think_path.is_file():
+            n_skipped_no_think += 1
+            continue
+        think_body = _read(think_path)
+        unified = f"{body}\n\n{THINK_MARKER}\n{think_body}\n"
+        if dry_run:
+            print(f"[dry-run] would unify {txt.name} ({len(body)} + {len(think_body)} chars)")
+        else:
+            txt.write_text(unified, encoding="utf-8")
+            think_path.unlink()
+        n_unified += 1
+    print(f"unified={n_unified} already-unified={n_skipped_already} no-think={n_skipped_no_think}")
+    return 0
+
+
+def cmd_pack_arrow(
+    folder: Path, out: Path, *,
+    image_extensions: tuple[str, ...] = (".jpg", ".jpeg", ".png", ".webp"),
+    inline_bytes: bool = True,
+) -> int:
+    """Pack a paired-folder dataset into a single parquet shard."""
+    try:
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+    except ImportError:
+        print("pyarrow required (pip install pyarrow)", file=sys.stderr); return 2
+
+    if not folder.is_dir():
+        print(f"folder not found: {folder}", file=sys.stderr); return 2
+
+    sample_ids: list[str] = []
+    captions: list[str] = []
+    thinks: list[str | None] = []
+    image_blobs: list[bytes | None] = []
+    image_paths: list[str | None] = []
+
+    pairs: list[tuple[Path, Path]] = []
+    for ext in image_extensions:
+        for img in sorted(folder.glob(f"*{ext}")):
+            txt = img.with_suffix(".txt")
+            if txt.is_file():
+                pairs.append((img, txt))
+    if not pairs:
+        print(f"no paired (image, .txt) pairs in {folder}", file=sys.stderr); return 2
+
+    for img_path, txt_path in pairs:
+        # Use the same parser used by PairedFolderT2IDataset
+        from train_u1.data.datasets import parse_caption_and_think
+        raw = _read(txt_path)
+        caption, think = parse_caption_and_think(raw)
+        # Legacy fallback for unmigrated data
+        if think is None:
+            tt = txt_path.with_suffix(".think.txt")
+            if tt.is_file():
+                think = _read(tt) or None
+        sample_ids.append(img_path.stem)
+        captions.append(caption)
+        thinks.append(think)
+        if inline_bytes:
+            with open(img_path, "rb") as f:
+                image_blobs.append(f.read())
+            image_paths.append(None)
+        else:
+            image_blobs.append(None)
+            image_paths.append(str(img_path.resolve()))
+
+    arrays = {
+        "sample_id": pa.array(sample_ids, type=pa.string()),
+        "caption": pa.array(captions, type=pa.string()),
+        "think": pa.array(thinks, type=pa.string()),
+    }
+    if inline_bytes:
+        arrays["image"] = pa.array(image_blobs, type=pa.binary())
+    else:
+        arrays["image_path"] = pa.array(image_paths, type=pa.string())
+
+    table = pa.table(arrays)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    pq.write_table(table, str(out), compression="zstd")
+    size_mb = out.stat().st_size / 1e6
+    print(f"wrote {table.num_rows} rows → {out}  ({size_mb:.1f} MB, "
+          f"inline_bytes={inline_bytes})")
+    return 0
+
+
+def cmd_inspect_arrow(path: Path) -> int:
+    try:
+        import pyarrow.parquet as pq
+    except ImportError:
+        print("pyarrow required", file=sys.stderr); return 2
+
+    table = pq.read_table(str(path))
+    print(f"path: {path}")
+    print(f"rows: {table.num_rows}")
+    print(f"schema:")
+    for f in table.schema:
+        print(f"  {f.name:<14s}  {f.type}")
+    print(f"---first 3 rows (truncated)---")
+    for i in range(min(3, table.num_rows)):
+        row = table.slice(i, 1).to_pydict()
+        sid = row["sample_id"][0]
+        cap = row["caption"][0][:80] + "..." if len(row["caption"][0]) > 80 else row["caption"][0]
+        thk = row.get("think", [None])[0]
+        thk_str = (thk[:80] + "...") if thk and len(thk) > 80 else (thk or "<none>")
+        if "image" in table.column_names:
+            sz = len(row["image"][0]) if row["image"][0] else 0
+            print(f"  [{i}] id={sid}  cap={cap!r}  image={sz/1e3:.1f} KB")
+        else:
+            print(f"  [{i}] id={sid}  cap={cap!r}  image_path={row['image_path'][0]}")
+        print(f"        think: {thk_str!r}")
+    return 0
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(description="Dataset format conversion utilities.")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    p_unify = sub.add_parser("unify-txt", help="fold .txt + .think.txt into one .txt")
+    p_unify.add_argument("folder", type=Path)
+    p_unify.add_argument("--dry-run", action="store_true")
+
+    p_pack = sub.add_parser("pack-arrow", help="pack folder → parquet shard")
+    p_pack.add_argument("folder", type=Path)
+    p_pack.add_argument("--out", type=Path, required=True)
+    p_pack.add_argument("--paths-only", action="store_true",
+                        help="store image_path instead of inline image bytes "
+                             "(smaller parquet, but reads still hit the filesystem)")
+
+    p_inspect = sub.add_parser("inspect-arrow", help="print row count + sample rows")
+    p_inspect.add_argument("path", type=Path)
+
+    args = ap.parse_args()
+    if args.cmd == "unify-txt":
+        return cmd_unify_txt(args.folder, dry_run=args.dry_run)
+    if args.cmd == "pack-arrow":
+        return cmd_pack_arrow(args.folder, args.out, inline_bytes=not args.paths_only)
+    if args.cmd == "inspect-arrow":
+        return cmd_inspect_arrow(args.path)
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

train_u1/scripts/sample_t2i_offload.py

@@ -304,7 +304,7 @@ def t2i_generate_offload(
             think_text = tokenizer.decode(think_token_ids, skip_special_tokens=False)
             if verbose:
                 tt = think_text.strip().replace("\n", " ")
-                print(f"[offload]   think generated +{cond_t_idx_extra} tokens: {tt[:160]!r}", flush=True)
+                print(f"[offload]   think generated +{cond_t_idx_extra} tokens (full): {tt!r}", flush=True)
         else:
             cond_kv, _ = model._t2i_prefix_forward(cond_ids, cond_idx, cond_attn)
         uncond_kv, _ = model._t2i_prefix_forward(uncond_ids, uncond_idx, uncond_attn)

train_u1/scripts/train_bf16_offload.py

@@ -44,7 +44,7 @@ import torch
 from train_u1.config import TrainRunConfig, load_train_config
 from train_u1.constants import MODEL_ID, MODEL_SHA
 from train_u1.data.collators import CollatorConfig, SenseNovaU1Collator, to_device
-from train_u1.data.datasets import PairedFolderT2IDataset
+from train_u1.data.datasets import ArrowT2IDataset, PairedFolderT2IDataset
 from train_u1.model.loader import _resolve_local_snapshot, load_neo_chat
 from train_u1.model.losses import fm_loss_x0
 from train_u1.model.lora import (
@@ -380,13 +380,23 @@ def main() -> int:
         trust_remote_code=True,
     )
 
-    # dataset + collator
-    ds = PairedFolderT2IDataset(
-        args.data_dir,
-        cap_max_pixels=args.cap_max_pixels,
-        snap_bucket=args.snap_bucket,
-    )
-    n_use = min(args.n_samples, len(ds))
+    # dataset + collator. Auto-dispatch on data_dir suffix:
+    #   .parquet     → ArrowT2IDataset (single shard or directory)
+    #   anything else, including a `.parquet`-less directory → PairedFolderT2IDataset
+    _data_path = Path(args.data_dir)
+    if _data_path.suffix == ".parquet" or (_data_path.is_dir() and any(_data_path.glob("*.parquet"))):
+        ds = ArrowT2IDataset(
+            _data_path,
+            cap_max_pixels=args.cap_max_pixels,
+            snap_bucket=args.snap_bucket,
+        )
+    else:
+        ds = PairedFolderT2IDataset(
+            args.data_dir,
+            cap_max_pixels=args.cap_max_pixels,
+            snap_bucket=args.snap_bucket,
+        )
+    n_use = len(ds) if args.n_samples is None else min(args.n_samples, len(ds))
     samples = [ds[i] for i in range(n_use)]
     collator = SenseNovaU1Collator(
         tok,

train_u1/tests/test_dataset_format.py

@@ -0,0 +1,131 @@
+"""Tests for the unified caption+think .txt format and arrow dataset."""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+from train_u1.data.datasets import THINK_DELIMITER_RE, parse_caption_and_think
+
+
+def test_parse_plain_caption() -> None:
+    c, t = parse_caption_and_think("A simple caption with no marker.")
+    assert c == "A simple caption with no marker."
+    assert t is None
+
+
+def test_parse_unified_format() -> None:
+    raw = (
+        "A painterly Hayateluc-style scene of a wisteria trellis.\n"
+        "\n"
+        "---think---\n"
+        "1. **Instruction Understanding:** The subject is a wisteria trellis...\n"
+        "6. **Explicit Prompt:** Render in Hayateluc's painterly anime style...\n"
+    )
+    c, t = parse_caption_and_think(raw)
+    assert c == "A painterly Hayateluc-style scene of a wisteria trellis."
+    assert t.startswith("1. **Instruction Understanding")
+    assert "6. **Explicit Prompt" in t
+
+
+def test_parse_marker_case_insensitive() -> None:
+    raw = "cap\n--- THINK ---\nthink body"
+    c, t = parse_caption_and_think(raw)
+    assert c == "cap"
+    assert t == "think body"
+
+
+def test_parse_extra_dashes() -> None:
+    raw = "cap\n----think----\nthink body"  # 4-dash variant
+    c, t = parse_caption_and_think(raw)
+    assert c == "cap"
+    assert t == "think body"
+
+
+def test_parse_no_marker_with_think_word() -> None:
+    """The literal word `think` (not on its own line, not as marker) is fine."""
+    raw = "A scene where the model needs to think hard about composition."
+    c, t = parse_caption_and_think(raw)
+    assert c == raw.strip()
+    assert t is None
+
+
+def test_parse_empty_think() -> None:
+    raw = "cap\n---think---\n   "
+    c, t = parse_caption_and_think(raw)
+    assert c == "cap"
+    assert t is None  # whitespace-only think → None
+
+
+def test_paired_folder_dataset_unified_txt(tmp_path: Path) -> None:
+    """PairedFolderT2IDataset reads the unified `<id>.txt` format."""
+    from PIL import Image
+
+    from train_u1.data.datasets import PairedFolderT2IDataset
+
+    img = Image.new("RGB", (64, 64), color=(128, 128, 128))
+    img.save(tmp_path / "test_001.jpg")
+    (tmp_path / "test_001.txt").write_text(
+        "An illustration by Hayateluc.\n---think---\nthink reasoning here\n"
+    )
+
+    ds = PairedFolderT2IDataset(tmp_path)
+    assert len(ds) == 1
+    s = ds[0]
+    assert s.sample_id == "test_001"
+    assert s.prompt == "An illustration by Hayateluc."
+    assert s.think == "think reasoning here"
+
+
+def test_paired_folder_dataset_legacy_split(tmp_path: Path) -> None:
+    """Legacy split format (`<id>.txt` + `<id>.think.txt`) still works."""
+    from PIL import Image
+
+    from train_u1.data.datasets import PairedFolderT2IDataset
+
+    img = Image.new("RGB", (64, 64), color=(0, 0, 0))
+    img.save(tmp_path / "split_001.jpg")
+    (tmp_path / "split_001.txt").write_text("plain caption only")
+    (tmp_path / "split_001.think.txt").write_text("legacy think text")
+
+    ds = PairedFolderT2IDataset(tmp_path)
+    s = ds[0]
+    assert s.prompt == "plain caption only"
+    assert s.think == "legacy think text"
+
+
+def test_arrow_dataset_roundtrip(tmp_path: Path) -> None:
+    """ArrowT2IDataset reads back what `dataset_tools pack-arrow` writes."""
+    pa = pytest.importorskip("pyarrow")
+    pq = pytest.importorskip("pyarrow.parquet")
+
+    from PIL import Image
+
+    from train_u1.data.datasets import ArrowT2IDataset
+    from train_u1.scripts.dataset_tools import cmd_pack_arrow
+
+    folder = tmp_path / "ds"
+    folder.mkdir()
+    img = Image.new("RGB", (64, 64), color=(255, 255, 255))
+    img.save(folder / "a.jpg")
+    (folder / "a.txt").write_text("cap a\n---think---\nthink a")
+
+    out = tmp_path / "out.parquet"
+    rc = cmd_pack_arrow(folder, out)
+    assert rc == 0
+    assert out.exists()
+
+    ds = ArrowT2IDataset(out)
+    assert len(ds) == 1
+    s = ds[0]
+    assert s.sample_id == "a"
+    assert s.prompt == "cap a"
+    assert s.think == "think a"
+    assert s.image.shape[0] == 3  # CHW
+
+
+def test_think_delimiter_regex_compiled() -> None:
+    """Sanity: regex matches the canonical marker."""
+    assert THINK_DELIMITER_RE.search("cap\n---think---\nthink") is not None
+    assert THINK_DELIMITER_RE.search("cap\n---tHiNk---\nthink") is not None
+    assert THINK_DELIMITER_RE.search("cap with think word") is None