ocr-foundry 0.2.0__py3-none-any.whl

khocr_gen/data_generator.py

@@ -0,0 +1,1160 @@
+"""DatasetGenerator - orchestrates synthetic data generation.
+
+Coordinates corpus loading, font selection, rendering with isolated augmentation, and parallel worker dispatch.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import os
+import random
+import shutil
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import cv2
+from tqdm import tqdm
+
+from .augmentation import write_output_image
+from .corpus import (
+    CorpusItem,
+    char_frequencies,
+    load_corpus_with_stats,
+    rare_chars_from_frequencies,
+)
+from .errors import InputValidationError
+from .fonts import FontManager
+from .normalizer import NormalizerConfig, normalize
+from .parallel import (
+    _init_render_worker,
+    _render_sample_batch,
+    derive_sample_seed,
+    ensure_worker_main_guard,
+    resolve_chunk_size,
+    resolve_mp_context,
+    resolve_worker_count,
+    seed_sample_generators,
+)
+from .provenance import build_sample_records, record_run_manifest
+from .rendering import ImageRenderer
+from .text_support import detect_script, line_separator_error
+
+if TYPE_CHECKING:
+    from .config import GenerationConfig
+
+_log = logging.getLogger("khocr_gen.generator")
+
+
+class DatasetGenerator:
+    """Generate synthetic training datasets from a text corpus.
+
+    Each text line is rendered to a clean canvas,
+    then exactly one augmentation method is applied per output image (isolated augmentation).
+    """
+
+    def __init__(self, cfg: GenerationConfig) -> None:
+        self._cfg = cfg
+        self.font_manager = FontManager(
+            language=cfg.language,
+            fonts_dir=cfg.fonts_dir,
+            multilingual=cfg.multilingual,
+            font_pool=cfg.font_pool,
+        )
+        self.renderer = ImageRenderer(self.font_manager, cfg)
+        self._run_ingestion_stats: dict[str, Any] = {}
+        self._run_used_fonts: dict[str, set[str]] = {}
+        self._used_font_paths: set[str] = set()
+
+    # Helpers
+
+    @staticmethod
+    def _reject_multilingual_line_separators(text_file: str | Path) -> None:
+        """Fail closed on U+2028/U+2029 before strip() can hide them.
+
+        The check uses the raw file records. Leading or trailing separators are
+        Unicode whitespace, so rejecting only the stripped line would keep the
+        visible remainder and expand it into copies.
+        """
+        rejected: list[str] = []
+        total = 0
+        with open(text_file, encoding="utf-8") as handle:
+            for raw in handle:
+                reason = line_separator_error(raw)
+                if reason is None and not raw.strip():
+                    continue
+                total += 1
+                if reason is None:
+                    continue
+                preview = raw.strip() or raw
+                rejected.append(f"{preview[:60]!r}: {reason}")
+        if rejected:
+            raise ValueError(
+                "Multilingual corpus preflight rejected "
+                f"{len(rejected)} / {total} lines containing a line or paragraph separator; "
+                "U+2028 and U+2029 are not horizontal-line text. " + "; ".join(rejected[:5])
+            )
+
+    @staticmethod
+    def _sanitize_label_text(text: str) -> str:
+        """Normalize embedded tabs in label text (labels.txt uses tab as delimiter)."""
+        return text.replace("\t", " ")
+
+    @staticmethod
+    def _read_non_empty_lines(
+        text_file: str | Path,
+        image_dir: str | None = None,
+        normalizer: NormalizerConfig | None = None,
+        stats: dict[str, int] | None = None,
+    ) -> list[str | CorpusItem | tuple[str, str]]:
+        """Read lines from *text_file*, optionally parsing (filename, text) pairs.
+
+        Auto-detects JSONL format by file extension, reconstituting
+        :class:`CorpusItem` instances with provenance metadata.
+        """
+        from .corpus import _is_jsonl, _parse_jsonl_record
+
+        text_file = Path(text_file)
+        lines: list[str | CorpusItem | tuple[str, str]] = []
+        if stats is None:
+            stats = {}
+        stats.update(total=0, passing=0, blank=0, empty=0, malformed=0)
+
+        if _is_jsonl(text_file) and not image_dir:
+            with open(text_file, encoding="utf-8-sig") as f:
+                for line_number, raw in enumerate(f, start=1):
+                    record = _parse_jsonl_record(raw, line_number, strict=True)
+                    if record is None:
+                        stats["blank"] += 1
+                        continue
+                    stats["total"] += 1
+                    text = record["text"]
+                    # Mirror load_corpus/_classify_raw_line: reject line
+                    # separators before strip() hides leading/trailing ones,
+                    # then skip blank/whitespace-only record text.
+                    reason = line_separator_error(text)
+                    if reason is not None:
+                        raise InputValidationError(
+                            f"JSONL line {line_number}: {reason}", code="jsonl_text_error"
+                        )
+                    text = text.strip()
+                    if not text:
+                        stats["empty"] += 1
+                        continue
+                    raw_val = record.get("raw_text", text)
+                    if "normalization_policy" in record and record.get("normalization_policy") is not None:
+                        norm_text = text
+                    elif normalizer is not None:
+                        norm_text = normalize(text, normalizer)
+                    else:
+                        norm_text = text
+                    norm_text = DatasetGenerator._sanitize_label_text(norm_text)
+                    if not norm_text:
+                        stats["empty"] += 1
+                        continue
+                    item = CorpusItem(
+                        norm_text,
+                        source_id=record.get("source_id"),
+                        source_group=record.get("source_group"),
+                        language=record.get("language"),
+                        script=record.get("script"),
+                        domain=record.get("domain"),
+                        source_uri=record.get("source_uri"),
+                        license=record.get("license"),
+                        retrieved_at=record.get("retrieved_at"),
+                        source_sha256=record.get("source_sha256"),
+                        normalization_policy=record.get("normalization_policy"),
+                        consent_or_restrictions=record.get("consent_or_restrictions"),
+                        raw_text=raw_val,
+                    )
+                    lines.append(item)
+                    stats["passing"] += 1
+            return lines
+
+        with open(text_file, encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    stats["blank"] += 1
+                    continue
+                stats["total"] += 1
+                if image_dir:
+                    parts = line.split("\t", 1)
+                    if len(parts) == 2:
+                        lines.append((parts[0], DatasetGenerator._sanitize_label_text(parts[1])))
+                        stats["passing"] += 1
+                    else:
+                        stats["malformed"] += 1
+                else:
+                    lines.append(DatasetGenerator._sanitize_label_text(line))
+                    stats["passing"] += 1
+        return lines
+
+    @staticmethod
+    def _split_lines_without_text_overlap(
+        lines: list, val_ratio: float = 0.1, test_ratio: float = 0.1, seed: int | None = 42
+    ) -> tuple[list, list, list]:
+        """Split *lines* into train/val/test keeping each group in exactly one split.
+
+        When items carry a ``source_group`` attribute (i.e. are
+        :class:`CorpusItem` instances loaded from a JSONL manifest), lines
+        are grouped by ``source_group`` so that all lines from the same
+        document/source stay together (**document-disjoint** splitting).
+        If no ``source_group`` is present, lines are grouped by their text
+        value (the existing **text-disjoint** behavior).
+
+        Returns ``(train_lines, val_lines, test_lines)`` where no group key
+        appears in more than one partition.
+
+        Resolution of effective ratios matches :meth:`GenerationConfig.resolve_split_ratios`:
+        pass ``test_ratio=0.0`` to disable the test split (returns empty list for test).
+        """
+        if not lines:
+            return [], [], []
+
+        val_ratio = float(max(0.0, min(1.0, val_ratio)))
+        test_ratio = float(max(0.0, min(1.0, test_ratio)))
+        # Clamp so train can't go negative
+        if val_ratio + test_ratio >= 1.0:
+            excess = val_ratio + test_ratio
+            val_ratio = val_ratio / excess
+            test_ratio = test_ratio / excess
+
+        if val_ratio <= 0.0 and test_ratio <= 0.0:
+            return list(lines), [], []
+
+        def _group_key(item: Any) -> str:
+            """Return the grouping key for a corpus item."""
+            # source_group takes priority (document-disjoint). Blank/whitespace-only
+            # groups fall back to text so they don't collapse into one bucket.
+            sg = getattr(item, "source_group", None)
+            if isinstance(sg, str) and sg.strip():
+                return sg
+            # Fall back to text value (text-disjoint, existing behavior)
+            return item[1] if isinstance(item, tuple) else item
+
+        buckets: dict[str, list] = {}
+        for item in lines:
+            key = _group_key(item)
+            buckets.setdefault(key, []).append(item)
+
+        unique_keys = list(buckets.keys())
+        if len(unique_keys) <= 1:
+            return list(lines), [], []
+
+        rng = random.Random(seed)
+        rng.shuffle(unique_keys)
+
+        n = len(unique_keys)
+        val_unique = round(n * val_ratio)
+        test_unique = round(n * test_ratio)
+        # Ensure at least one key stays in train
+        val_unique = max(0, min(val_unique, n - 1))
+        test_unique = max(0, min(test_unique, n - 1 - val_unique))
+
+        val_keys = set(unique_keys[:val_unique])
+        test_keys = set(unique_keys[val_unique : val_unique + test_unique])
+
+        train_lines: list = []
+        val_lines: list = []
+        test_lines: list = []
+        for key in unique_keys:
+            if key in val_keys:
+                val_lines.extend(buckets[key])
+            elif key in test_keys:
+                test_lines.extend(buckets[key])
+            else:
+                train_lines.extend(buckets[key])
+
+        return train_lines, val_lines, test_lines
+
+    @staticmethod
+    def _copies_for_line(
+        text: str,
+        copies: int,
+        rare_chars: set[str] | None,
+        rare_char_multiplier: float,
+    ) -> int:
+        """Boost *copies* for lines containing a rare character.
+
+        Lines with no rare character are unaffected. Multiplier is only
+        applied (and only ever increases copies) when rendering random-font
+        training samples; val/test always pass ``rare_chars=None``.
+        """
+        if not rare_chars or rare_char_multiplier <= 1.0:
+            return copies
+        if any(ch in rare_chars for ch in text):
+            return max(copies, round(copies * rare_char_multiplier))
+        return copies
+
+    @staticmethod
+    def _write_lines(path: Path, lines: list) -> None:
+        with open(path, "w", encoding="utf-8") as f:
+            for line in lines:
+                if isinstance(line, CorpusItem):
+                    record: dict[str, Any] = {"text": str(line)}
+                    for attr in (
+                        "source_id",
+                        "source_group",
+                        "language",
+                        "script",
+                        "domain",
+                        "source_uri",
+                        "license",
+                        "retrieved_at",
+                        "source_sha256",
+                        "normalization_policy",
+                        "consent_or_restrictions",
+                        "raw_text",
+                    ):
+                        val = getattr(line, attr, None)
+                        if val is not None:
+                            record[attr] = val
+                    f.write(json.dumps(record, ensure_ascii=False) + "\n")
+                elif isinstance(line, tuple):
+                    f.write(f"{line[0]}\t{line[1]}\n")
+                else:
+                    f.write(line + "\n")
+
+    @staticmethod
+    def _build_vocab(labels_files: list[str], output_path: str) -> None:
+        """Scan labels files and write a vocab.json covering every character."""
+        print("\nBuilding vocabulary from generated labels...")
+        unique_chars: set[str] = set()
+
+        for lf in labels_files:
+            if not os.path.exists(lf):
+                continue
+            with open(lf, encoding="utf-8") as fh:
+                for line in fh:
+                    parts = line.strip().split("\t", 1)
+                    if len(parts) == 2:
+                        unique_chars.update(parts[1])
+
+        print(f"   Found {len(unique_chars)} unique characters.")
+
+        if not unique_chars:
+            print(
+                "   WARNING: no characters found in any labels file. "
+                "The written vocab.json will contain only {'<unk>': 0}."
+            )
+
+        vocab: dict[str, int] = {"<unk>": 0}
+        for idx, ch in enumerate(sorted(unique_chars), start=1):
+            if ch != "<unk>":
+                vocab[ch] = idx
+
+        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+        with open(output_path, "w", encoding="utf-8") as fh:
+            json.dump(vocab, fh, indent=2, ensure_ascii=False)
+
+        print(f"   Vocabulary saved -> {output_path}")
+
+    # Worker config
+
+    def _render_worker_config(
+        self,
+        output_dir: Path,
+        split_name: str,
+        font_mode: str,
+        retry_limit: int,
+        image_dir: str | None = None,
+    ) -> dict[str, Any]:
+        """Build the dict passed to worker initializer."""
+        worker_dict = self._cfg.to_dict()
+        worker_dict.update(
+            {
+                "split_name": split_name,
+                "output_images_dir": str(output_dir / "images"),
+                "font_mode": font_mode,
+                "retry_limit": retry_limit,
+                "image_dir": image_dir,
+            }
+        )
+        return worker_dict
+
+    # Batch iteration
+
+    @staticmethod
+    def _iter_sample_batches(samples: list[tuple[str, Any]], start_idx: int, chunk_size: int):
+        """Yield batches of (index, text, font_ref) tuples."""
+        batch: list[tuple[int, str, Any]] = []
+        for offset, (text, font_ref) in enumerate(samples):
+            batch.append((start_idx + offset, text, font_ref))
+            if len(batch) >= chunk_size:
+                yield batch
+                batch = []
+        if batch:
+            yield batch
+
+    # Split generation
+
+    def _generate_split_serial(
+        self,
+        samples: list[tuple[str, Any]],
+        output_dir: Path,
+        split_name: str,
+        font_mode: str,
+        retry_limit: int,
+        start_idx: int,
+        file_mode: str,
+        image_dir: str | None = None,
+    ) -> tuple[int, list[int]]:
+        """Generate one split using a single process.
+
+        Returns `(success_count, heights)` where *heights* is the pixel height
+        of every successfully rendered image (for the QA height-distribution
+        summary).
+        """
+        enabled_methods = self._cfg.enabled_aug_methods()
+        success_count = 0
+        heights: list[int] = []
+        record_metadata = bool(self._cfg.record_metadata)
+        # Provenance is only meaningful for JSONL manifests, which yield
+        # CorpusItem instances carrying source identity. Plain .txt corpora
+        # and image-dir mode produce no provenance.jsonl.
+        has_corpus_items = any(isinstance(text, CorpusItem) for text, _ in samples)
+        prov_rows: list[dict[str, Any]] = []
+
+        meta_file = (
+            open(output_dir / "metadata.jsonl", file_mode, encoding="utf-8")  # noqa: SIM115
+            if record_metadata
+            else None
+        )
+        try:
+            with open(output_dir / "labels.txt", file_mode, encoding="utf-8") as labels_file:
+                for idx, (text, specific_font) in enumerate(
+                    tqdm(samples, desc="    Generating", unit="img")
+                ):
+                    current_idx = start_idx + idx
+                    sample_id = f"{split_name}_{current_idx:06d}"
+                    seed_sample_generators(self._cfg.seed, current_idx, sample_id)
+                    current_retry_limit = 1 if font_mode == "all" else retry_limit
+
+                    try:
+                        meta: dict[str, Any] = {}
+                        if image_dir:
+                            img_path = Path(image_dir) / specific_font
+                            img_array = cv2.imread(str(img_path))
+                            if img_array is None:
+                                raise RuntimeError(f"Could not read image: {img_path}")
+                            if img_array.ndim == 3 and img_array.shape[2] == 3:
+                                img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
+
+                            aug_results = self.renderer.augment_image(
+                                img_array, enabled_methods, script=detect_script(text)
+                            )
+                            if not aug_results:
+                                continue
+                            _, img = aug_results[0]
+                        else:
+                            result = self.renderer.render_with_one_augmentation(
+                                text,
+                                enabled_methods,
+                                specific_font=specific_font,
+                                retry_limit=current_retry_limit,
+                            )
+                            if result is None:
+                                continue
+                            _, img, meta = result
+
+                        if img is None:
+                            continue
+
+
+                        img_filename = f"{split_name}_{current_idx:06d}.{self._cfg.output_format}"
+                        img_path = output_dir / "images" / img_filename
+                        write_ok = write_output_image(
+                            img_path,
+                            img,
+                            fmt=self._cfg.output_format,
+                            jpeg_quality=self._cfg.jpeg_quality,
+                        )
+                        if not write_ok:
+                            raise RuntimeError(f"cv2.imwrite failed for output path: {img_path}")
+
+                        labels_file.write(f"{img_filename}\t{text}\n")
+                        heights.append(int(img.shape[0]))
+                        self._used_font_paths.update(
+                            path for path in meta.get("fonts", [meta.get("font")]) if path
+                        )
+                        if meta_file is not None:
+                            record, prov_row = build_sample_records(
+                                img_filename=img_filename,
+                                text=text,
+                                width=img.shape[1],
+                                height=img.shape[0],
+                                split_name=split_name,
+                                meta=meta,
+                                default_language=self._cfg.language,
+                            )
+                            if prov_row is not None:
+                                prov_rows.append(prov_row)
+                            meta_file.write(json.dumps(record, ensure_ascii=False) + "\n")
+                        success_count += 1
+
+                    except Exception as e:
+                        print(f"    Failed for '{text[:30]}...': {e}")
+                        continue
+        finally:
+            if meta_file is not None:
+                meta_file.close()
+
+        if record_metadata and has_corpus_items:
+            prov_path = output_dir / "provenance.jsonl"
+            with open(prov_path, file_mode, encoding="utf-8") as prov_file:
+                for row in prov_rows:
+                    prov_file.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+        return success_count, heights
+
+    def _generate_split_parallel(
+        self,
+        samples: list[tuple[str, Any]],
+        output_dir: Path,
+        split_name: str,
+        font_mode: str,
+        retry_limit: int,
+        start_idx: int,
+        file_mode: str,
+        workers: int,
+        image_dir: str | None = None,
+    ) -> tuple[int, list[int]]:
+        """Generate one split using multiple worker processes.
+
+        Returns `(success_count, heights)`; see `_generate_split_serial`.
+        """
+        import json as _json
+        from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, TimeoutError, wait
+
+        resolved_workers = resolve_worker_count(workers, len(samples))
+        if resolved_workers <= 1:
+            return self._generate_split_serial(
+                samples=samples,
+                output_dir=output_dir,
+                split_name=split_name,
+                font_mode=font_mode,
+                retry_limit=retry_limit,
+                start_idx=start_idx,
+                file_mode=file_mode,
+                image_dir=image_dir,
+            )
+
+        chunk_size = resolve_chunk_size(len(samples), resolved_workers)
+        if (
+            resolved_workers > 1
+            and len(samples) > resolved_workers
+            and chunk_size > len(samples) // resolved_workers
+        ):
+            chunk_size = max(1, len(samples) // resolved_workers)
+        worker_config = self._render_worker_config(
+            output_dir=output_dir,
+            split_name=split_name,
+            font_mode=font_mode,
+            retry_limit=retry_limit,
+            image_dir=image_dir,
+        )
+
+        print(f"  Rendering with {resolved_workers} worker processes (chunk size: {chunk_size})")
+
+        worker_timeout = self._cfg.worker_timeout if self._cfg.worker_timeout > 0 else None
+        error_log_path = output_dir / "generation_errors.jsonl"
+        record_metadata = bool(self._cfg.record_metadata)
+        # Provenance is only meaningful for JSONL manifests, which yield
+        # CorpusItem instances carrying source identity.
+        has_corpus_items = any(isinstance(text, CorpusItem) for text, _ in samples)
+
+        success_count = 0
+        heights: list[int] = []
+        error_records: list[dict] = []
+        batch_iter = self._iter_sample_batches(samples, start_idx=start_idx, chunk_size=chunk_size)
+        max_in_flight = max(2, resolved_workers * 2)
+
+        mp_context, start_method = resolve_mp_context()
+        ensure_worker_main_guard(start_method)
+
+        with (
+            open(output_dir / "labels.txt", file_mode, encoding="utf-8") as labels_file,
+            contextlib.ExitStack() as stack,
+            ProcessPoolExecutor(
+                max_workers=resolved_workers,
+                mp_context=mp_context,
+                initializer=_init_render_worker,
+                initargs=(worker_config,),
+            ) as executor,
+        ):
+            meta_file = (
+                stack.enter_context(
+                    open(output_dir / "metadata.jsonl", file_mode, encoding="utf-8")
+                )
+                if record_metadata
+                else None
+            )
+            prov_file = (
+                stack.enter_context(
+                    open(output_dir / "provenance.jsonl", file_mode, encoding="utf-8")
+                )
+                if record_metadata and has_corpus_items
+                else None
+            )
+            print(f"  Multiprocessing start method: {start_method}")
+            in_flight: dict[Any, int] = {}
+            pending_results: dict[int, Any] = {}
+            batch_seq = 0
+            next_write_seq = 0
+
+            for _ in range(max_in_flight):
+                try:
+                    batch = next(batch_iter)
+                except StopIteration:
+                    break
+                fut = executor.submit(_render_sample_batch, batch)
+                in_flight[fut] = batch_seq
+                batch_seq += 1
+
+            with tqdm(total=len(samples), desc="    Generating", unit="img") as pbar:
+                while in_flight:
+                    done, _ = wait(set(in_flight.keys()), return_when=FIRST_COMPLETED)
+                    for future in done:
+                        b_seq = in_flight.pop(future)
+                        try:
+                            res = future.result(timeout=worker_timeout)
+                        except TimeoutError:
+                            raise RuntimeError(
+                                f"A worker batch timed out after {worker_timeout}s. "
+                                "Increase --worker-timeout or set it to 0 to disable."
+                            ) from None
+                        pending_results[b_seq] = res
+                        pbar.update(res[1])
+
+                        try:
+                            batch = next(batch_iter)
+                            new_fut = executor.submit(_render_sample_batch, batch)
+                            in_flight[new_fut] = batch_seq
+                            batch_seq += 1
+                        except StopIteration:
+                            pass
+
+                    while next_write_seq in pending_results:
+                        (
+                            labels,
+                            _attempted,
+                            batch_success,
+                            err_msgs,
+                            meta_lines,
+                            prov_lines,
+                            batch_heights,
+                        ) = pending_results.pop(next_write_seq)
+                        next_write_seq += 1
+
+                        if labels:
+                            labels_file.writelines(labels)
+                        if meta_file is not None and meta_lines:
+                            meta_file.writelines(meta_lines)
+                        for meta_line in meta_lines:
+                            worker_meta = json.loads(meta_line)
+                            self._used_font_paths.update(worker_meta.get("fonts", []))
+                        if prov_file is not None and prov_lines:
+                            prov_file.writelines(prov_lines)
+                        success_count += batch_success
+                        heights.extend(batch_heights)
+
+                        for message in err_msgs:
+                            print(message)
+                            error_records.append({"message": message})
+        if error_records:
+            with open(error_log_path, "a", encoding="utf-8") as ef:
+                for rec in error_records:
+                    ef.write(_json.dumps(rec, ensure_ascii=False) + "\n")
+            print(
+                f"  {len(error_records)} error(s) written to "
+                f"{error_log_path.relative_to(output_dir.parent)}"
+            )
+
+        return success_count, heights
+
+    # Top-level entrypoints
+
+    def generate_split(
+        self,
+        text_file: str | Path,
+        output_dir: str | Path,
+        split_name: str,
+        font_mode: str = "random",
+        retry_limit: int = 10,
+        append: bool = False,
+        workers: int = 0,
+        image_dir: str | None = None,
+        copies: int = 3,
+        rare_chars: set[str] | None = None,
+        rare_char_multiplier: float = 1.0,
+    ) -> int:
+        """Generate one split (train or val) from *text_file*.
+
+        *rare_chars*/*rare_char_multiplier* boost the per-line copy count for
+        lines containing a rare character; pass ``rare_chars=None`` (the
+        default) for splits that should stay at a uniform ``copies`` per line
+        (e.g. val/test).
+        """
+        output_dir = Path(output_dir)
+        (output_dir / "images").mkdir(parents=True, exist_ok=True)
+        overrides = {
+            "text_file": str(text_file), "font_mode": font_mode, "copies": copies,
+            "retry_limit": retry_limit, "workers": workers, "append": append,
+            "image_dir": image_dir, "rare_chars": sorted(rare_chars or ()),
+            "rare_char_multiplier": rare_char_multiplier,
+        }
+
+        if self._cfg.multilingual and not image_dir:
+            self._reject_multilingual_line_separators(text_file)
+
+        ingestion_stats: dict[str, int] = {}
+        lines = self._read_non_empty_lines(
+            text_file, image_dir=image_dir, normalizer=self._cfg.normalizer, stats=ingestion_stats
+        )
+        self._run_ingestion_stats[split_name] = {
+            "path": str(text_file),
+            "counts": ingestion_stats,
+            "length_filters_applied": False,
+        }
+        self._used_font_paths = set()
+        print(f"  Loaded {len(lines)} lines from {text_file}")
+
+        if self._cfg.multilingual and not image_dir:
+            rejected: list[str] = []
+            for line in cast("list[str]", lines):
+                if not self.renderer.is_line_eligible(line):
+                    rejected.append(f"{line[:60]!r}: {self.renderer.last_rejection}")
+            if rejected:
+                req_desc = (
+                    "provide fonts covering all grapheme clusters or select another --font-pool"
+                    if self._cfg.cluster_font_fallback
+                    else "provide a font covering each whole line or select another --font-pool"
+                )
+                raise ValueError(
+                    f"Multilingual font preflight rejected {len(rejected)} / {len(lines)} lines; "
+                    f"{req_desc}. "
+                    + "; ".join(rejected[:5])
+                )
+
+        # Build samples
+        samples: list[tuple[str, Any]] = []
+
+        if image_dir:
+            print("  Using EXISTING IMAGES mode (retry limit ignored)")
+            for item in lines:
+                filename, text = item
+                n = self._copies_for_line(text, copies, rare_chars, rare_char_multiplier)
+                for _ in range(n):
+                    samples.append((text, filename))
+        elif font_mode == "all":
+            print("  Using ALL fonts mode")
+            fonts_list = self.font_manager.all_fonts
+            print(f"  Iterating {len(fonts_list)} fonts per line...")
+            # image_dir is falsy here, so _read_non_empty_lines only produced plain str lines.
+            for line in cast("list[str]", lines):
+                if self._cfg.multilingual:
+                    eligible = set(self.renderer.eligible_font_paths(line))
+                    if eligible:
+                        for font_tuple in fonts_list:
+                            if font_tuple[0] in eligible:
+                                samples.append((line, (font_tuple[0], font_tuple[1])))
+                    elif self._cfg.cluster_font_fallback and self.renderer.is_line_eligible(line):
+                        samples.append((line, None))
+                else:
+                    for font_tuple in fonts_list:
+                        samples.append((line, (font_tuple[0], font_tuple[1])))
+        else:
+            print(f"  Using RANDOM fonts mode (retry limit: {retry_limit} attempts per sample)")
+            for line in cast("list[str]", lines):
+                n = self._copies_for_line(line, copies, rare_chars, rare_char_multiplier)
+                for _ in range(n):
+                    samples.append((line, None))
+
+        # samples now holds every (text, font_ref) entry to render; the raw
+        # line list isn't needed anymore and would otherwise stay resident
+        # alongside samples for the rest of the split.
+        del lines
+
+        if not samples:
+            print("  Generating 0 images...")
+            self._record_split_run(output_dir, split_name, overrides, attempted=0, generated=0)
+            return 0
+
+        if self._cfg.seed is not None:
+            shuffle_seed = derive_sample_seed(self._cfg.seed, 0, f"split_shuffle_{split_name}")
+            random.Random(shuffle_seed).shuffle(samples)
+        else:
+            random.shuffle(samples)
+        file_mode = "a" if append else "w"
+
+        print(f"  Generating {len(samples)} images...")
+
+        success_count, heights = self._generate_split_parallel(
+            samples=samples,
+            output_dir=output_dir,
+            split_name=split_name,
+            font_mode=font_mode,
+            retry_limit=retry_limit,
+            start_idx=0,
+            file_mode=file_mode,
+            workers=workers,
+            image_dir=image_dir,
+        )
+
+        self._record_split_run(
+            output_dir, split_name, overrides, attempted=len(samples), generated=success_count,
+        )
+        print(f"  Generated {success_count} / {len(samples)} images")
+        if self._cfg.multilingual:
+            print(f"  Failed/skipped: {len(samples) - success_count}")
+            if success_count == 0:
+                raise ValueError(
+                    "Multilingual generation produced zero images; check font loading, RAQM, and augmentation settings"
+                )
+        self._print_height_summary(heights)
+        print()
+        return success_count
+
+    def _record_split_run(
+        self, output_dir: Path, split_name: str, overrides: dict[str, Any],
+        *, attempted: int, generated: int,
+    ) -> None:
+        self._run_used_fonts[split_name] = set(self._used_font_paths)
+        record_run_manifest(
+            output_dir,
+            config=self._cfg,
+            split_name=split_name,
+            invocation_overrides=overrides,
+            available_fonts=self.font_manager.get_available_fonts(),
+            used_font_paths=sorted(self._used_font_paths),
+            ingestion={split_name: self._run_ingestion_stats[split_name]},
+            outcomes={"attempted": attempted, "generated": generated, "failed_or_skipped": attempted - generated},
+        )
+
+    def _print_height_summary(self, heights: list[int]) -> None:
+        """Print observed line-height distribution for QA (variable-height modes)."""
+        mode = self._cfg.line_height_mode
+        print(f"  Line height mode: {mode}")
+        if mode != "fixed":
+            print(
+                f"  Height range: {self._cfg.min_line_height}..{self._cfg.max_line_height} "
+                f"step {self._cfg.line_height_step}"
+            )
+        if heights:
+            mean = sum(heights) / len(heights)
+            print(f"  Observed heights: min={min(heights)} max={max(heights)} mean={mean:.1f}")
+
+    def generate_dataset(
+        self,
+        corpus_path: str | Path,
+        output_dir: str | Path = "data",
+        val_file: str | Path | None = None,
+        test_file: str | Path | None = None,
+        copies: int = 3,
+        font_mode: str = "random",
+        retry_limit: int = 10,
+        existing_mode: str = "error",
+        workers: int = 0,
+        split_seed: int | None = 42,
+        val_ratio: float = 0.1,
+        test_ratio: float = 0.1,
+        image_dir: str | None = None,
+        min_length: int = 5,
+        max_length: int = 120,
+        max_lines: int = 0,
+        oversample_rare_chars: bool = False,
+        rare_char_percentile: float = 5.0,
+        rare_char_multiplier: float = 3.0,
+    ) -> dict[str, int]:
+        """Generate a complete dataset with train/val/test splits.
+
+        When *oversample_rare_chars* is set, training lines containing one of
+        the least-frequent *rare_char_percentile* percent of characters in
+        the corpus get their copy count multiplied by *rare_char_multiplier*
+        (val/test stay uniform, so evaluation reflects natural frequency).
+
+        Returns a dict of split -> sample count.
+        """
+        if split_seed is None and self._cfg.seed is not None:
+            split_seed = self._cfg.seed
+        valid_modes = {"error", "append", "overwrite"}
+        if existing_mode not in valid_modes:
+            raise ValueError(
+                f"existing_mode must be one of {sorted(valid_modes)}, got: {existing_mode!r}"
+            )
+        if not 0.0 <= float(val_ratio) < 1.0:
+            raise ValueError("val_ratio must be in [0.0, 1.0).")
+        if not 0.0 <= float(test_ratio) < 1.0:
+            raise ValueError("test_ratio must be in [0.0, 1.0).")
+
+        val_ratio = float(val_ratio)
+        test_ratio = float(test_ratio)
+
+        # An explicit test_file is the sole source of the test split; don't also
+        # carve a test set out of the disjoint ratio-based split (mirrors val_file).
+        if test_file and Path(test_file).exists():
+            test_ratio = 0.0
+
+        print("\n" + "=" * 70)
+        print("  Dataset Generation")
+        print("=" * 70)
+
+        output_dir = Path(output_dir)
+
+        if self._cfg.multilingual and not image_dir:
+            self._reject_multilingual_line_separators(corpus_path)
+            for extra in (val_file, test_file):
+                if extra and Path(extra).exists():
+                    self._reject_multilingual_line_separators(extra)
+
+        # Handle existing output
+        append_mode = False
+        if output_dir.exists() and any(output_dir.iterdir()):
+            print(f"\n  Dataset directory '{output_dir}' already exists.")
+            if existing_mode == "append":
+                append_mode = True
+                print("  Continuing generation (appending to existing)...")
+            elif existing_mode == "overwrite":
+                print("  Cleaning up existing directory...")
+                shutil.rmtree(output_dir)
+            else:
+                raise FileExistsError(
+                    "Dataset output already exists. Re-run with "
+                    "existing_mode='append' or existing_mode='overwrite'."
+                )
+
+        # Create output dirs
+        output_dir.mkdir(exist_ok=True)
+        (output_dir / "train" / "images").mkdir(parents=True, exist_ok=True)
+
+        will_have_val = bool((val_file and Path(val_file).exists()) or float(val_ratio) > 0.0)
+        if will_have_val:
+            (output_dir / "val" / "images").mkdir(parents=True, exist_ok=True)
+
+        will_have_test = bool((test_file and Path(test_file).exists()) or float(test_ratio) > 0.0)
+        if will_have_test:
+            (output_dir / "test" / "images").mkdir(parents=True, exist_ok=True)
+
+        counts: dict[str, int] = {}
+        self._run_ingestion_stats = {}
+        self._run_used_fonts = {}
+
+        rare_chars: set[str] | None = None
+        if oversample_rare_chars:
+            freq = char_frequencies(
+                corpus_path,
+                min_length=min_length,
+                max_length=max_length,
+                max_lines=max_lines,
+                normalizer=self._cfg.normalizer,
+            )
+            rare_chars = rare_chars_from_frequencies(freq, percentile=rare_char_percentile)
+            print(
+                f"  Rare-char oversampling: {len(rare_chars)} rare char(s) "
+                f"(bottom {rare_char_percentile:g}% of {len(freq)} distinct chars), "
+                f"{rare_char_multiplier:g}x copies for lines containing one (train only)"
+            )
+
+        if val_file and Path(val_file).exists():
+            print("\nGenerating TRAINING set...")
+            counts["train"] = self.generate_split(
+                text_file=corpus_path,
+                output_dir=output_dir / "train",
+                split_name="train",
+                font_mode=font_mode,
+                retry_limit=retry_limit,
+                append=append_mode,
+                workers=workers,
+                image_dir=image_dir,
+                copies=copies,
+                rare_chars=rare_chars,
+                rare_char_multiplier=rare_char_multiplier,
+            )
+
+            print("\nGenerating VALIDATION set (from separate file)...")
+            counts["val"] = self.generate_split(
+                text_file=val_file,
+                output_dir=output_dir / "val",
+                split_name="val",
+                font_mode=font_mode,
+                retry_limit=retry_limit,
+                append=append_mode,
+                workers=workers,
+                image_dir=image_dir,
+                copies=1,
+            )
+
+            if test_file and Path(test_file).exists():
+                print("\nGenerating TEST set (from separate file)...")
+                counts["test"] = self.generate_split(
+                    text_file=test_file,
+                    output_dir=output_dir / "test",
+                    split_name="test",
+                    font_mode=font_mode,
+                    retry_limit=retry_limit,
+                    append=append_mode,
+                    workers=workers,
+                    image_dir=image_dir,
+                    copies=1,
+                )
+        else:
+            print("\nPreparing disjoint TRAIN/VAL/TEST split from corpus...")
+            print(
+                f"  Split config: train={(1.0 - val_ratio - test_ratio) * 100:.1f}%, "
+                f"val={val_ratio * 100:.1f}%, "
+                f"test={test_ratio * 100:.1f}%, seed={split_seed}"
+            )
+
+            # Filter the corpus straight into a single in-memory list (one
+            # materialization) instead of round-tripping through a temp file
+            # and re-reading it back into a second list.
+            filtered_lines, corpus_stats = load_corpus_with_stats(
+                corpus_path,
+                min_length=min_length,
+                max_length=max_length,
+                max_lines=max_lines,
+                normalizer=self._cfg.normalizer,
+                strict_jsonl=True,
+            )
+            self._run_ingestion_stats["corpus"] = {
+                "path": str(corpus_path),
+                "counts": corpus_stats,
+                "length_filters_applied": True,
+                "min_length": min_length, "max_length": max_length, "max_lines": max_lines,
+                "count_scope": "records inspected until the passing-record cap; unread remainder is not counted",
+            }
+            all_lines = list(tqdm(filtered_lines, desc="  Filtering corpus", unit="line"))
+            train_lines, val_lines, test_lines = self._split_lines_without_text_overlap(
+                all_lines, val_ratio=val_ratio, test_ratio=test_ratio, seed=split_seed
+            )
+            # all_lines is now fully partitioned across train/val/test; drop the
+            # reference so it doesn't sit in memory alongside the per-split
+            # re-read that generate_split performs below.
+            del all_lines
+
+            train_unique = len({(t[1] if isinstance(t, tuple) else t) for t in train_lines})
+            val_unique = len({(t[1] if isinstance(t, tuple) else t) for t in val_lines})
+            test_unique = len({(t[1] if isinstance(t, tuple) else t) for t in test_lines})
+
+            def _text_of(item):
+                return item[1] if isinstance(item, tuple) else item
+
+            overlap_tv = len({_text_of(x) for x in train_lines} & {_text_of(x) for x in val_lines})
+            overlap_tt = len({_text_of(x) for x in train_lines} & {_text_of(x) for x in test_lines})
+            overlap_vt = len({_text_of(x) for x in val_lines} & {_text_of(x) for x in test_lines})
+            print(
+                f"  Split summary: train={len(train_lines)} lines ({train_unique} unique), "
+                f"val={len(val_lines)} lines ({val_unique} unique), "
+                f"test={len(test_lines)} lines ({test_unique} unique), "
+                f"overlaps: train∩val={overlap_tv} train∩test={overlap_tt} val∩test={overlap_vt}"
+            )
+
+            temp_ext = ".jsonl" if train_lines and isinstance(train_lines[0], CorpusItem) else ".txt"
+            temp_train = output_dir / f"temp_train{temp_ext}"
+            temp_val = output_dir / f"temp_val{temp_ext}"
+            temp_test = output_dir / f"temp_test{temp_ext}"
+            self._write_lines(temp_train, train_lines)
+            self._write_lines(temp_val, val_lines)
+            self._write_lines(temp_test, test_lines)
+            # The split lists are now on disk for generate_split to stream back
+            # in per-split. Drop them so the full corpus isn't held twice while
+            # rendering runs.
+            val_lines_present = bool(val_lines)
+            test_lines_present = bool(test_lines)
+            del train_lines, val_lines, test_lines
+
+            try:
+                print("\nGenerating TRAINING set...")
+                counts["train"] = self.generate_split(
+                    text_file=temp_train,
+                    output_dir=output_dir / "train",
+                    split_name="train",
+                    font_mode=font_mode,
+                    retry_limit=retry_limit,
+                    append=append_mode,
+                    workers=workers,
+                    image_dir=image_dir,
+                    copies=copies,
+                    rare_chars=rare_chars,
+                    rare_char_multiplier=rare_char_multiplier,
+                )
+
+                if val_lines_present:
+                    print("\nGenerating VALIDATION set (disjoint text split)...")
+                    counts["val"] = self.generate_split(
+                        text_file=temp_val,
+                        output_dir=output_dir / "val",
+                        split_name="val",
+                        font_mode=font_mode,
+                        retry_limit=retry_limit,
+                        append=append_mode,
+                        workers=workers,
+                        image_dir=image_dir,
+                        copies=1,
+                    )
+
+                if test_lines_present:
+                    print("\nGenerating TEST set (disjoint text split)...")
+                    counts["test"] = self.generate_split(
+                        text_file=temp_test,
+                        output_dir=output_dir / "test",
+                        split_name="test",
+                        font_mode=font_mode,
+                        retry_limit=retry_limit,
+                        append=append_mode,
+                        workers=workers,
+                        image_dir=image_dir,
+                        copies=1,
+                    )
+
+                # Handle separate test_file when val_file path was NOT the primary path
+                if test_file and Path(test_file).exists():
+                    print("\nGenerating TEST set (from separate file)...")
+                    counts["test"] = self.generate_split(
+                        text_file=test_file,
+                        output_dir=output_dir / "test",
+                        split_name="test",
+                        font_mode=font_mode,
+                        retry_limit=retry_limit,
+                        append=append_mode,
+                        workers=workers,
+                        image_dir=image_dir,
+                        copies=1,
+                    )
+            finally:
+                for tmp_path in (temp_train, temp_val, temp_test):
+                    with contextlib.suppress(FileNotFoundError):
+                        tmp_path.unlink()
+
+        record_run_manifest(
+            output_dir,
+            config=self._cfg,
+            invocation_overrides={
+                "corpus_path": str(corpus_path),
+                "val_file": str(val_file) if val_file else None,
+                "test_file": str(test_file) if test_file else None,
+                "copies": copies, "font_mode": font_mode, "retry_limit": retry_limit,
+                "existing_mode": existing_mode, "workers": workers, "split_seed": split_seed,
+                "val_ratio": val_ratio, "test_ratio": test_ratio, "image_dir": image_dir,
+                "min_length": min_length, "max_length": max_length, "max_lines": max_lines,
+                "oversample_rare_chars": oversample_rare_chars,
+                "rare_char_percentile": rare_char_percentile,
+                "rare_char_multiplier": rare_char_multiplier,
+                "split_grouping": "explicit_files" if val_file else "source_group_or_exact_text",
+            },
+            available_fonts=self.font_manager.get_available_fonts(),
+            used_font_paths=sorted(set().union(*self._run_used_fonts.values())),
+            ingestion=self._run_ingestion_stats,
+            outcomes={"generated_by_split": counts},
+        )
+        print("\n" + "=" * 70)
+        print("  Dataset Generation Complete!")
+        print("=" * 70)
+        for split, count in counts.items():
+            print(f"  {split.capitalize()}: {count:,} samples")
+        print(f"  Output: {output_dir}")
+        print("=" * 70 + "\n")
+
+        return counts