REPO: D:\mintdim
GENERATED: 2026-05-29 10:45:07
PURPOSE: inspect CLI command adapter and unit-build UI logger separation.

====================================================================================================
FILE: src\mintdim\cli\main.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: import sys
    4: 
    5: from mintdim.cli.commands import unit_build
    6: 
    7: 
    8: def main(argv: list[str] | None = None) -> int:
    9:     args = list(sys.argv[1:] if argv is None else argv)
   10: 
   11:     if not args:
   12:         _print_usage()
   13:         return 2
   14: 
   15:     command = args.pop(0)
   16: 
   17:     if command == "unit-build":
   18:         return unit_build.main(args)
   19: 
   20:     sys.stderr.write(f"Unknown command: {command}\n")
   21:     _print_usage()
   22:     return 2
   23: 
   24: 
   25: def _print_usage() -> None:
   26:     sys.stderr.write("Usage: mintdim <command> [args...]\n")
   27:     sys.stderr.write("Commands:\n")
   28:     sys.stderr.write("  unit-build\n")
   29: 
   30: 
   31: if __name__ == "__main__":
   32:     raise SystemExit(main())

====================================================================================================
FILE: src\mintdim\cli\commands\unit_build.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: from mintdim.apis.unit_build.api import UnitBuildPipeline
    4: from mintdim.apis.unit_build.ui.logger import main as ui_main
    5: 
    6: 
    7: def create_pipeline() -> UnitBuildPipeline:
    8:     return UnitBuildPipeline()
    9: 
   10: 
   11: def main(argv: list[str] | None = None) -> int:
   12:     return ui_main(argv)

====================================================================================================
FILE: src\mintdim\apis\unit_build\ui\logger.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: import os
    4: import shutil
    5: import sys
    6: import time
    7: from dataclasses import dataclass
    8: from typing import Any, Mapping, TextIO
    9: 
   10: for _stream in (sys.stdout, sys.stderr):
   11:     if _stream.encoding and _stream.encoding.lower() != "utf-8":
   12:         _stream.reconfigure(encoding="utf-8")
   13: 
   14: APP_NAME = "unit-build"
   15: 
   16: 
   17: def _detect_path_mark() -> str:
   18:     try:
   19:         "⎿".encode(sys.stderr.encoding or "utf-8")
   20:         return "⎿"
   21:     except (UnicodeEncodeError, LookupError):
   22:         return "->"
   23: 
   24: 
   25: PATH_MARK = _detect_path_mark()
   26: 
   27: SPINNER_FRAMES = ("-", "\\", "|", "/")
   28: 
   29: BAR_WIDTH_DEFAULT = 22
   30: BAR_WIDTH_MIN = 8
   31: 
   32: NON_TTY_PROGRESS_INTERVAL = 1.0
   33: 
   34: _progress_active = False
   35: _last_progress_at = 0.0
   36: _progress_last_fields: dict[str, Any] | None = None
   37: 
   38: 
   39: @dataclass(frozen=True)
   40: class RGB:
   41:     r: int
   42:     g: int
   43:     b: int
   44: 
   45: 
   46: class Palette:
   47:     green = RGB(34, 197, 94)
   48:     green_soft = RGB(74, 222, 128)
   49:     blue = RGB(59, 130, 246)
   50:     cyan = RGB(34, 211, 238)
   51:     red = RGB(220, 38, 38)
   52:     text = RGB(226, 232, 240)
   53:     muted = RGB(100, 116, 139)
   54:     dim = RGB(148, 163, 184)
   55:     shard_bg = RGB(22, 32, 46)
   56: 
   57: 
   58: def main(argv: list[str] | None = None) -> int:
   59:     """Minimal unit-build CLI entrypoint.
   60: 
   61:     The public Python API is still the main supported interface. This function
   62:     exists so the package-level dispatcher can import and route to unit-build
   63:     without failing.
   64:     """
   65:     args = list(sys.argv[1:] if argv is None else argv)
   66: 
   67:     if not args or args[0] in {"-h", "--help", "help"}:
   68:         sys.stderr.write("Usage: mintdim unit-build [options]\n")
   69:         sys.stderr.write("\n")
   70:         sys.stderr.write("This command currently exposes the unit-build logger/UI.\n")
   71:         sys.stderr.write("Use the Python API for full dataset builds.\n")
   72:         return 0
   73: 
   74:     sys.stderr.write(f"Unknown unit-build argument: {args[0]}\n")
   75:     return 2
   76: 
   77: 
   78: def log_event(event: str, **fields: Any) -> None:
   79:     """
   80:     Stream-friendly stderr logger.
   81: 
   82:     Important:
   83:     - progress is indeterminate, no percent and no ETA
   84:     - progress is kept short to avoid terminal wrapping
   85:     - output paths are intentionally not shown in progress lines
   86:     - UI text is ASCII-safe for Windows terminals and redirected logs
   87: 
   88:     Layout policy:
   89:     - there is only one live progress line
   90:     - shard events are durable log lines above the live progress line
   91:     - the previous progress line is cleared in place before any non-progress log
   92:     """
   93:     global _progress_last_fields
   94: 
   95:     if event == "progress":
   96:         _render_progress(fields)
   97:         return
   98: 
   99:     if event in ("start", "done", "failed"):
  100:         _finish_progress_line()
  101:         _progress_last_fields = None
  102: 
  103:         if event == "start":
  104:             _render_start(fields)
  105:         elif event == "done":
  106:             _render_done(fields)
  107:         else:
  108:             _render_failed(fields)
  109: 
  110:         return
  111: 
  112:     was_active = _progress_active and sys.stderr.isatty()
  113:     _clear_progress_line()
  114: 
  115:     if event == "shard":
  116:         _render_shard(fields)
  117:     else:
  118:         _render_event(event, fields)
  119: 
  120:     if was_active and _progress_last_fields is not None:
  121:         _render_progress(_progress_last_fields)
  122: 
  123: 
  124: def log_result(result: Mapping[str, Any]) -> None:
  125:     """
  126:     Replacement for:
  127:         print("Done:", result)
  128: 
  129:     Prints a compact final summary instead of a raw Python dict.
  130:     """
  131:     _finish_progress_line()
  132: 
  133:     rows: list[tuple[str, Any]] = []
  134: 
  135:     output = (
  136:         result.get("output")
  137:         or result.get("output_dir")
  138:         or _format_outputs(result.get("outputs"))
  139:     )
  140: 
  141:     samples = result.get("samples", result.get("total_samples"))
  142:     tokens = result.get("tokens", result.get("total_tokens"))
  143:     shards = result.get("shards", result.get("total_shards"))
  144: 
  145:     if output is not None:
  146:         rows.append(("output", output))
  147: 
  148:     if samples is not None:
  149:         rows.append(("samples", _compact_number(samples)))
  150: 
  151:     if tokens is not None:
  152:         rows.append(("tokens", _compact_number(tokens)))
  153: 
  154:     if shards is not None:
  155:         rows.append(("shards", shards))
  156: 
  157:     for key in ("samples_with_unk", "duplicate_groups"):
  158:         value = result.get(key)
  159:         if value is not None:
  160:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  161: 
  162:     _render_card(
  163:         icon="OK",
  164:         title="complete",
  165:         accent=Palette.green,
  166:         rows=rows or [("status", "ok")],
  167:     )
  168: 
  169: 
  170: def _render_start(fields: dict[str, Any]) -> None:
  171:     rows: list[tuple[str, Any]] = []
  172: 
  173:     output = fields.get("output", fields.get("output_dir"))
  174:     units = fields.get("units")
  175:     batch = fields.get("build_batch", fields.get("batch"))
  176:     files = fields.get("files")
  177: 
  178:     if output is not None:
  179:         rows.append(("output", output))
  180: 
  181:     if units is not None:
  182:         rows.append(("units", _format_units(units)))
  183: 
  184:     if batch is not None:
  185:         rows.append(("batch", batch))
  186: 
  187:     if files is not None:
  188:         rows.append(("files", files))
  189: 
  190:     used = {"output", "output_dir", "units", "build_batch", "batch", "files"}
  191: 
  192:     for key, value in fields.items():
  193:         if key not in used:
  194:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  195: 
  196:     _render_card(
  197:         icon=">",
  198:         title="started",
  199:         accent=Palette.blue,
  200:         rows=rows,
  201:     )
  202: 
  203: 
  204: def _render_progress(fields: dict[str, Any]) -> None:
  205:     global _progress_active, _last_progress_at, _progress_last_fields
  206: 
  207:     now = time.monotonic()
  208:     stream = sys.stderr
  209: 
  210:     _progress_last_fields = dict(fields)
  211: 
  212:     if not stream.isatty():
  213:         if now - _last_progress_at < NON_TTY_PROGRESS_INTERVAL:
  214:             return
  215: 
  216:         _last_progress_at = now
  217:         _render_progress_snapshot(fields)
  218:         return
  219: 
  220:     spinner = SPINNER_FRAMES[int(now * 10) % len(SPINNER_FRAMES)]
  221:     width = _bar_width()
  222:     bar_plain = "#" * width
  223:     bar_colored = _stream_bar(width=width, tick=now)
  224: 
  225:     metrics = _progress_metrics(fields)
  226: 
  227:     plain = f"{spinner} {APP_NAME} {bar_plain}"
  228:     colored = (
  229:         f"{_fg(Palette.green)}{spinner}{_reset()} "
  230:         f"{_strong(APP_NAME)} "
  231:         f"{bar_colored}"
  232:     )
  233: 
  234:     for label, value in metrics:
  235:         piece_plain = f"  {label} {value}"
  236:         piece_colored = f"  {_label(label)} {_value(value)}"
  237: 
  238:         candidate_plain = plain + piece_plain
  239: 
  240:         if len(candidate_plain) > _terminal_width() - 1:
  241:             break
  242: 
  243:         plain = candidate_plain
  244:         colored += piece_colored
  245: 
  246:     print(f"\r\x1b[2K{colored}{_reset()}", end="", file=stream, flush=True)
  247: 
  248:     _progress_active = True
  249:     _last_progress_at = now
  250: 
  251: 
  252: def _render_progress_snapshot(fields: dict[str, Any]) -> None:
  253:     metrics = _progress_metrics(fields)
  254: 
  255:     if not metrics:
  256:         return
  257: 
  258:     print(
  259:         f"{APP_NAME} "
  260:         + " ".join(f"{label}={value}" for label, value in metrics),
  261:         file=sys.stderr,
  262:         flush=True,
  263:     )
  264: 
  265: 
  266: def _render_done(fields: dict[str, Any]) -> None:
  267:     rows: list[tuple[str, Any]] = []
  268: 
  269:     output = fields.get("output", fields.get("output_dir"))
  270:     samples = fields.get("samples", fields.get("total_samples"))
  271:     tokens = fields.get("tokens", fields.get("total_tokens"))
  272:     shards = fields.get("shards", fields.get("total_shards"))
  273: 
  274:     if output is not None:
  275:         rows.append(("output", output))
  276: 
  277:     if samples is not None:
  278:         rows.append(("samples", _compact_number(samples)))
  279: 
  280:     if tokens is not None:
  281:         rows.append(("tokens", _compact_number(tokens)))
  282: 
  283:     if shards is not None:
  284:         rows.append(("shards", shards))
  285: 
  286:     for key in ("samples_with_unk", "duplicate_groups"):
  287:         value = fields.get(key)
  288:         if value is not None:
  289:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  290: 
  291:     used = {
  292:         "output",
  293:         "output_dir",
  294:         "samples",
  295:         "total_samples",
  296:         "tokens",
  297:         "total_tokens",
  298:         "shards",
  299:         "total_shards",
  300:         "samples_with_unk",
  301:         "duplicate_groups",
  302:     }
  303: 
  304:     for key, value in fields.items():
  305:         if key not in used:
  306:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  307: 
  308:     _render_card(
  309:         icon="OK",
  310:         title="complete",
  311:         accent=Palette.green,
  312:         rows=rows or [("status", "ok")],
  313:     )
  314: 
  315: 
  316: def _render_failed(fields: dict[str, Any]) -> None:
  317:     rows: list[tuple[str, Any]] = []
  318: 
  319:     preferred = (
  320:         "output",
  321:         "output_dir",
  322:         "batches",
  323:         "samples",
  324:         "total_samples",
  325:         "tokens",
  326:         "total_tokens",
  327:         "shards",
  328:         "total_shards",
  329:         "error",
  330:     )
  331: 
  332:     used: set[str] = set()
  333: 
  334:     for key in preferred:
  335:         value = fields.get(key)
  336:         if value is not None:
  337:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  338:             used.add(key)
  339: 
  340:     for key, value in fields.items():
  341:         if key not in used:
  342:             rows.append((_pretty_label(key), _format_value_for_key(key, value)))
  343: 
  344:     _render_card(
  345:         icon="X",
  346:         title="failed",
  347:         accent=Palette.red,
  348:         rows=rows or [("error", "unknown")],
  349:     )
  350: 
  351: 
  352: def _render_shard(fields: dict[str, Any]) -> None:
  353:     metrics: list[tuple[str, Any]] = []
  354: 
  355:     path = fields.get("path")
  356: 
  357:     preferred = (
  358:         "shard",
  359:         "samples",
  360:         "total_samples",
  361:         "tokens",
  362:         "total_tokens",
  363:         "bytes",
  364:         "size",
  365:     )
  366: 
  367:     used: set[str] = {"path"}
  368: 
  369:     for key in preferred:
  370:         value = fields.get(key)
  371: 
  372:         if value is None:
  373:             continue
  374: 
  375:         label = _pretty_label(key)
  376: 
  377:         if label == "samples" and any(item[0] == "samples" for item in metrics):
  378:             used.add(key)
  379:             continue
  380: 
  381:         if label == "tokens" and any(item[0] == "tokens" for item in metrics):
  382:             used.add(key)
  383:             continue
  384: 
  385:         metrics.append((label, _format_value_for_key(key, value)))
  386:         used.add(key)
  387: 
  388:     for key, value in fields.items():
  389:         if key not in used:
  390:             metrics.append((_pretty_label(key), _format_value_for_key(key, value)))
  391: 
  392:     _emit_line_with_bg(
  393:         sys.stderr,
  394:         f"{_fg(Palette.cyan)}>{_reset()} "
  395:         f"{_strong(APP_NAME)} "
  396:         f"{_fg(Palette.cyan)}shard{_reset()}  "
  397:         f"{_inline_metrics(metrics)}",
  398:         bg=Palette.shard_bg,
  399:     )
  400: 
  401:     if path is not None:
  402:         _emit_line(
  403:             sys.stderr,
  404:             f"  {_fg(Palette.dim)}{PATH_MARK}{_reset()} "
  405:             f"{_label('Path:')} {_value(path)}",
  406:         )
  407: 
  408: 
  409: def _render_event(event: str, fields: dict[str, Any]) -> None:
  410:     metrics = [
  411:         (_pretty_label(key), _format_value_for_key(key, value))
  412:         for key, value in fields.items()
  413:     ]
  414: 
  415:     if metrics:
  416:         line = (
  417:             f"{_fg(Palette.dim)}>{_reset()} "
  418:             f"{_strong(APP_NAME)} "
  419:             f"{_fg(Palette.cyan)}{event}{_reset()}  "
  420:             f"{_inline_metrics(metrics)}"
  421:         )
  422:     else:
  423:         line = (
  424:             f"{_fg(Palette.dim)}>{_reset()} "
  425:             f"{_strong(APP_NAME)} "
  426:             f"{_fg(Palette.cyan)}{event}{_reset()}"
  427:         )
  428: 
  429:     _emit_line(sys.stderr, line)
  430: 
  431: 
  432: def _render_card(
  433:     *,
  434:     icon: str,
  435:     title: str,
  436:     accent: RGB,
  437:     rows: list[tuple[str, Any]],
  438: ) -> None:
  439:     stream = sys.stderr
  440:     label_width = max((len(str(label)) for label, _ in rows), default=0)
  441: 
  442:     _emit_line(
  443:         stream,
  444:         f"{_fg(accent)}{icon}{_reset()} "
  445:         f"{_strong(APP_NAME)} "
  446:         f"{_fg(accent)}{title}{_reset()}",
  447:     )
  448: 
  449:     for label, value in rows:
  450:         _emit_line(
  451:             stream,
  452:             f"  {_label(f'{label:<{label_width}}')}  {_value(value)}",
  453:         )
  454: 
  455: 
  456: def _emit_line(stream: TextIO, line: str) -> None:
  457:     width = _terminal_width()
  458: 
  459:     if stream.isatty():
  460:         print(
  461:             f"\x1b[2K{_truncate_ansi(line, width - 1)}{_reset()}",
  462:             file=stream,
  463:             flush=True,
  464:         )
  465:     else:
  466:         print(_strip_ansi(_truncate_ansi(line, width - 1)), file=stream, flush=True)
  467: 
  468: 
  469: def _emit_line_with_bg(stream: TextIO, inner: str, bg: RGB) -> None:
  470:     """Render `inner` with a flat background color spanning the full row.
  471: 
  472:     The inner string may use `_reset()` (`\\x1b[0m`) for fg/bold styling; we
  473:     rewrite those to a partial reset (`\\x1b[22;39m`) so the bg attribute is
  474:     not torn down mid-line. The line is padded with spaces up to `width - 1`
  475:     so the bg color fills the entire row before the final full reset.
  476:     """
  477:     if not (stream.isatty() and _supports_color()):
  478:         _emit_line(stream, inner)
  479:         return
  480: 
  481:     width = _terminal_width()
  482:     target = width - 1
  483: 
  484:     inner_kept = inner.replace("\x1b[0m", "\x1b[22;39m")
  485:     visible = len(_strip_ansi(inner_kept))
  486: 
  487:     if visible >= target:
  488:         body = _truncate_ansi(inner_kept, target)
  489:     else:
  490:         body = inner_kept + " " * (target - visible)
  491: 
  492:     print(f"\x1b[2K{_bg(bg)}{body}\x1b[0m", file=stream, flush=True)
  493: 
  494: 
  495: def _finish_progress_line() -> None:
  496:     global _progress_active
  497: 
  498:     if not _progress_active:
  499:         return
  500: 
  501:     print(_reset(), file=sys.stderr, flush=True)
  502:     _progress_active = False
  503: 
  504: 
  505: def _clear_progress_line() -> None:
  506:     """Wipe the progress line in place without leaving a leftover row.
  507: 
  508:     Used when a mid-stream event, such as a shard rollover, needs to be printed
  509:     above the live progress line. Cursor returns to column 0 of the same row so
  510:     the caller's print writes the event there and scrolls it up; the progress
  511:     bar is then redrawn on the new bottom row.
  512: 
  513:     This function is intentionally not `_finish_progress_line()`: it must not
  514:     preserve the old progress line as history.
  515:     """
  516:     global _progress_active
  517: 
  518:     if not _progress_active:
  519:         return
  520: 
  521:     if sys.stderr.isatty():
  522:         print(f"\r\x1b[2K{_reset()}", end="", file=sys.stderr, flush=True)
  523: 
  524:     _progress_active = False
  525: 
  526: 
  527: def _progress_metrics(fields: dict[str, Any]) -> list[tuple[str, Any]]:
  528:     """
  529:     Keep live progress stable and short.
  530: 
  531:     The progress line is a live status line, not a full log record. It should
  532:     not expose paths or secondary counters because those make the terminal wrap
  533:     and cause metrics to disappear as numbers grow.
  534: 
  535:     Durable details belong in shard/done/failed events.
  536:     """
  537:     preferred = (
  538:         "samples",
  539:         "total_samples",
  540:         "batch",
  541:         "tokens",
  542:         "total_tokens",
  543:     )
  544: 
  545:     used: set[str] = set()
  546:     metrics: list[tuple[str, Any]] = []
  547: 
  548:     for key in preferred:
  549:         value = fields.get(key)
  550: 
  551:         if value is None:
  552:             continue
  553: 
  554:         label = _pretty_label(key)
  555: 
  556:         if label == "samples" and any(item[0] == "samples" for item in metrics):
  557:             used.add(key)
  558:             continue
  559: 
  560:         if label == "tokens" and any(item[0] == "tokens" for item in metrics):
  561:             used.add(key)
  562:             continue
  563: 
  564:         metrics.append((label, _format_value_for_key(key, value)))
  565:         used.add(key)
  566: 
  567:     return metrics
  568: 
  569: 
  570: def _inline_metrics(metrics: list[tuple[str, Any]]) -> str:
  571:     return "  |  ".join(
  572:         f"{_label(label)} {_value(value)}"
  573:         for label, value in metrics
  574:         if value is not None
  575:     )
  576: 
  577: 
  578: def _stream_bar(*, width: int, tick: float) -> str:
  579:     sweep_width = max(4, width // 4)
  580:     position = int(tick * 14) % (width + sweep_width)
  581:     start = position - sweep_width
  582: 
  583:     chars: list[str] = []
  584: 
  585:     for index in range(width):
  586:         active = start <= index < start + sweep_width
  587: 
  588:         if active:
  589:             chars.append(f"{_fg(Palette.green_soft)}#")
  590:         else:
  591:             chars.append(f"{_fg(Palette.muted)}-")
  592: 
  593:     return "".join(chars) + _reset()
  594: 
  595: 
  596: def _bar_width() -> int:
  597:     terminal_width = _terminal_width()
  598: 
  599:     if terminal_width < 70:
  600:         return BAR_WIDTH_MIN
  601: 
  602:     if terminal_width < 95:
  603:         return 14
  604: 
  605:     return BAR_WIDTH_DEFAULT
  606: 
  607: 
  608: def _terminal_width() -> int:
  609:     return shutil.get_terminal_size(fallback=(88, 24)).columns
  610: 
  611: 
  612: def _pretty_label(label: str) -> str:
  613:     aliases = {
  614:         "bytes": "size",
  615:         "total_samples": "samples",
  616:         "total_tokens": "tokens",
  617:         "batch_samples": "+",
  618:         "build_batch": "batch",
  619:         "output_dir": "output",
  620:         "samples_with_unk": "unknown",
  621:         "duplicate_groups": "duplicates",
  622:         "total_shards": "shards",
  623:     }
  624: 
  625:     return aliases.get(label, label.replace("_", " "))
  626: 
  627: 
  628: def _format_units(units: Any) -> str:
  629:     if isinstance(units, (list, tuple, set)):
  630:         return ", ".join(str(unit) for unit in units)
  631: 
  632:     return str(units)
  633: 
  634: 
  635: def _format_outputs(outputs: Any) -> str | None:
  636:     if not outputs:
  637:         return None
  638: 
  639:     if isinstance(outputs, (list, tuple)):
  640:         if len(outputs) == 1:
  641:             return str(outputs[0])
  642: 
  643:         return f"{len(outputs)} files"
  644: 
  645:     return str(outputs)
  646: 
  647: 
  648: def _format_value_for_key(key: str, value: Any) -> Any:
  649:     if key in {
  650:         "samples",
  651:         "total_samples",
  652:         "tokens",
  653:         "total_tokens",
  654:         "batch_tokens",
  655:         "samples_with_unk",
  656:     }:
  657:         return _compact_number(value)
  658: 
  659:     if key in {"bytes", "size"}:
  660:         return _compact_size(value)
  661: 
  662:     return value
  663: 
  664: 
  665: def _compact_number(value: Any) -> Any:
  666:     if not isinstance(value, int):
  667:         return value
  668: 
  669:     abs_value = abs(value)
  670: 
  671:     if abs_value >= 1_000_000_000:
  672:         return f"{value / 1_000_000_000:.2f}B"
  673: 
  674:     if abs_value >= 1_000_000:
  675:         return f"{value / 1_000_000:.2f}M"
  676: 
  677:     if abs_value >= 10_000:
  678:         return f"{value / 1_000:.1f}K"
  679: 
  680:     return str(value)
  681: 
  682: 
  683: def _compact_size(value: Any) -> Any:
  684:     if not isinstance(value, int):
  685:         return value
  686: 
  687:     abs_value = abs(value)
  688: 
  689:     if abs_value >= 1024**3:
  690:         return f"{value / 1024**3:.2f}GB"
  691: 
  692:     if abs_value >= 1024**2:
  693:         return f"{value / 1024**2:.2f}MB"
  694: 
  695:     if abs_value >= 1024:
  696:         return f"{value / 1024:.1f}KB"
  697: 
  698:     return f"{value}B"
  699: 
  700: 
  701: def _supports_color() -> bool:
  702:     if os.getenv("NO_COLOR"):
  703:         return False
  704: 
  705:     if os.getenv("FORCE_COLOR"):
  706:         return True
  707: 
  708:     return sys.stderr.isatty()
  709: 
  710: 
  711: def _fg(rgb: RGB) -> str:
  712:     if not _supports_color():
  713:         return ""
  714: 
  715:     return f"\x1b[38;2;{rgb.r};{rgb.g};{rgb.b}m"
  716: 
  717: 
  718: def _bg(rgb: RGB) -> str:
  719:     if not _supports_color():
  720:         return ""
  721: 
  722:     return f"\x1b[48;2;{rgb.r};{rgb.g};{rgb.b}m"
  723: 
  724: 
  725: def _reset() -> str:
  726:     if not _supports_color():
  727:         return ""
  728: 
  729:     return "\x1b[0m"
  730: 
  731: 
  732: def _strong(value: object) -> str:
  733:     if not _supports_color():
  734:         return str(value)
  735: 
  736:     return f"\x1b[1m{_fg(Palette.text)}{value}{_reset()}"
  737: 
  738: 
  739: def _label(value: object) -> str:
  740:     return f"{_fg(Palette.muted)}{value}{_reset()}"
  741: 
  742: 
  743: def _value(value: object) -> str:
  744:     return f"{_fg(Palette.green)}{value}{_reset()}"
  745: 
  746: 
  747: def _dim(value: object) -> str:
  748:     return f"{_fg(Palette.dim)}{value}{_reset()}"
  749: 
  750: 
  751: def _strip_ansi(text: str) -> str:
  752:     result: list[str] = []
  753:     index = 0
  754: 
  755:     while index < len(text):
  756:         char = text[index]
  757: 
  758:         if char == "\x1b":
  759:             index += 1
  760: 
  761:             if index < len(text) and text[index] == "[":
  762:                 index += 1
  763: 
  764:                 while index < len(text) and not text[index].isalpha():
  765:                     index += 1
  766: 
  767:                 if index < len(text):
  768:                     index += 1
  769: 
  770:             continue
  771: 
  772:         result.append(char)
  773:         index += 1
  774: 
  775:     return "".join(result)
  776: 
  777: 
  778: def _truncate_ansi(text: str, max_visible: int) -> str:
  779:     if max_visible <= 0:
  780:         return ""
  781: 
  782:     visible = 0
  783:     result: list[str] = []
  784:     index = 0
  785: 
  786:     while index < len(text):
  787:         char = text[index]
  788: 
  789:         if char == "\x1b":
  790:             start = index
  791:             index += 1
  792: 
  793:             if index < len(text) and text[index] == "[":
  794:                 index += 1
  795: 
  796:                 while index < len(text) and not text[index].isalpha():
  797:                     index += 1
  798: 
  799:                 if index < len(text):
  800:                     index += 1
  801: 
  802:                 result.append(text[start:index])
  803: 
  804:             continue
  805: 
  806:         if visible >= max_visible:
  807:             result.append("...")
  808:             break
  809: 
  810:         result.append(char)
  811:         visible += 1
  812:         index += 1
  813: 
  814:     return "".join(result)
  815: 
  816: 
  817: if __name__ == "__main__":
  818:     raise SystemExit(main())

====================================================================================================
FILE: src\mintdim\apis\unit_build\ui\prompts.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: import sys
    4: from typing import Literal, TextIO
    5: 
    6: from mintdim.apis.unit_build.ui.logger import (
    7:     Palette,
    8:     _emit_line,
    9:     _fg,
   10:     _label,
   11:     _reset,
   12:     _strong,
   13:     _value,
   14: )
   15: from mintdim.apis.unit_build.domain.overflow import (
   16:     OverflowAction,
   17:     OverflowContext,
   18:     OverflowResolution,
   19: )
   20: 
   21: 
   22: def stdin_is_tty() -> bool:
   23:     return sys.stdin.isatty()
   24: 
   25: 
   26: def _write_block(stream: TextIO, lines: list[str]) -> None:
   27:     for line in lines:
   28:         _emit_line(stream, line)
   29: 
   30: 
   31: def _read_choice(
   32:     stream_in: TextIO,
   33:     valid: set[str],
   34:     prompt: str,
   35: ) -> str:
   36:     while True:
   37:         sys.stderr.write(prompt)
   38:         sys.stderr.flush()
   39: 
   40:         raw = stream_in.readline()
   41: 
   42:         if not raw:
   43:             raise EOFError("stdin closed during prompt")
   44: 
   45:         choice = raw.strip().lower()
   46: 
   47:         if choice in valid:
   48:             return choice
   49: 
   50:         sys.stderr.write(
   51:             f"  invalid choice {choice!r}, expected one of {sorted(valid)}\n"
   52:         )
   53: 
   54: 
   55: def prompt_template_unk(
   56:     *,
   57:     file_index: int,
   58:     n_findings: int,
   59:     stream_in: TextIO | None = None,
   60: ) -> bool:
   61:     """Ask whether to continue when literal template segments produce UNK tokens.
   62: 
   63:     Returns True if the user types `y` or `yes`.
   64:     Returns False for Enter, `n`, `no`, or EOF.
   65:     """
   66:     inp = stream_in or sys.stdin
   67:     out = sys.stderr
   68:     title = "[unit-build] Template UNK detected"
   69: 
   70:     _write_block(
   71:         out,
   72:         [
   73:             f"{_fg(Palette.red)}{_strong(title)}{_reset()}",
   74:             f"  {_label('file index')} {_value(file_index)}",
   75:             f"  {_label('literal segments with UNK')} {_value(n_findings)}",
   76:             "",
   77:             "  Static-text tokens will be UNK in every encoded sample.",
   78:             "  Continue building?",
   79:             "",
   80:         ],
   81:     )
   82: 
   83:     while True:
   84:         sys.stderr.write("  [y/N] > ")
   85:         sys.stderr.flush()
   86: 
   87:         raw = inp.readline()
   88: 
   89:         if not raw:
   90:             return False
   91: 
   92:         ans = raw.strip().lower()
   93: 
   94:         if ans in ("", "n", "no"):
   95:             return False
   96: 
   97:         if ans in ("y", "yes"):
   98:             return True
   99: 
  100:         sys.stderr.write("  invalid, expected y or n\n")
  101: 
  102: 
  103: def prompt_overflow_action(
  104:     *,
  105:     ctx: OverflowContext,
  106:     stream_in: TextIO | None = None,
  107: ) -> OverflowAction:
  108:     """Ask how to handle a sample that exceeds max unit_size."""
  109:     inp = stream_in or sys.stdin
  110:     out = sys.stderr
  111:     title = "[unit-build] Sample exceeds max unit_size"
  112: 
  113:     _write_block(
  114:         out,
  115:         [
  116:             f"{_fg(Palette.red)}{_strong(title)}{_reset()}",
  117:             f"  {_label('source')}      {_value(f'{ctx.source_file}:{ctx.source_line}')}",
  118:             f"  {_label('sample_id')}   {_value(ctx.sample_id)}",
  119:             f"  {_label('token_count')} {_value(ctx.token_count)}",
  120:             f"  {_label('max(sizes)')}  {_value(ctx.max_unit_size)}",
  121:             "",
  122:             "  How should this sample be handled?",
  123:             "      [1] truncate  drop trailing tokens to fit max(sizes)",
  124:             "      [2] extend    declare a new unit_size that fits this sample",
  125:             "      [3] skip      drop this sample and record it in overflow_index.jsonl",
  126:             "",
  127:         ],
  128:     )
  129: 
  130:     choice = _read_choice(inp, {"1", "2", "3"}, "  [1/2/3] > ")
  131: 
  132:     return {
  133:         "1": OverflowAction.TRUNCATE,
  134:         "2": OverflowAction.EXTEND,
  135:         "3": OverflowAction.SKIP,
  136:     }[choice]
  137: 
  138: 
  139: def prompt_overflow_extend_size(
  140:     *,
  141:     token_count: int,
  142:     stream_in: TextIO | None = None,
  143: ) -> int:
  144:     """Read a new unit_size after the user chooses EXTEND."""
  145:     inp = stream_in or sys.stdin
  146:     out = sys.stderr
  147: 
  148:     _write_block(
  149:         out,
  150:         [
  151:             f"  {_label('new unit_size')} must be an integer >= {token_count}",
  152:         ],
  153:     )
  154: 
  155:     while True:
  156:         sys.stderr.write("  > ")
  157:         sys.stderr.flush()
  158: 
  159:         raw = inp.readline()
  160: 
  161:         if not raw:
  162:             raise EOFError("stdin closed during prompt")
  163: 
  164:         try:
  165:             value = int(raw.strip())
  166:         except ValueError:
  167:             sys.stderr.write("  invalid, expected an integer\n")
  168:             continue
  169: 
  170:         if value < token_count:
  171:             sys.stderr.write(f"  invalid, must be >= {token_count}\n")
  172:             continue
  173: 
  174:         return value
  175: 
  176: 
  177: def prompt_overflow_scope(
  178:     *,
  179:     stream_in: TextIO | None = None,
  180: ) -> Literal["only", "all"]:
  181:     """Ask whether the selected overflow action applies once or globally."""
  182:     inp = stream_in or sys.stdin
  183:     out = sys.stderr
  184: 
  185:     _write_block(
  186:         out,
  187:         [
  188:             "  Apply this choice to:",
  189:             "      [1] only this sample",
  190:             "      [2] all subsequent overflow samples in this run",
  191:         ],
  192:     )
  193: 
  194:     choice = _read_choice(inp, {"1", "2"}, "  [1/2] > ")
  195: 
  196:     return "only" if choice == "1" else "all"
  197: 
  198: 
  199: def resolve_overflow_via_prompt(
  200:     *,
  201:     ctx: OverflowContext,
  202:     stream_in: TextIO | None = None,
  203: ) -> tuple[OverflowResolution, Literal["only", "all"]]:
  204:     """Run the overflow prompt flow and return the selected resolution."""
  205:     action = prompt_overflow_action(ctx=ctx, stream_in=stream_in)
  206: 
  207:     new_size: int | None = None
  208: 
  209:     if action == OverflowAction.EXTEND:
  210:         new_size = prompt_overflow_extend_size(
  211:             token_count=ctx.token_count,
  212:             stream_in=stream_in,
  213:         )
  214: 
  215:     scope = prompt_overflow_scope(stream_in=stream_in)
  216: 
  217:     return (
  218:         OverflowResolution(action=action, new_unit_size=new_size),
  219:         scope,
  220:     )

====================================================================================================
FILE: src\mintdim\apis\unit_build\api.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: from dataclasses import dataclass, field
    4: from typing import Sequence
    5: 
    6: 
    7: @dataclass
    8: class SourceAPI:
    9:     pipeline: "UnitBuildPipeline"
   10: 
   11:     def jsonl(
   12:         self,
   13:         *,
   14:         files: Sequence[str],
   15:         fields: Sequence[Sequence[str]],
   16:         templates: Sequence[str],
   17:     ) -> "UnitBuildPipeline":
   18:         self.pipeline._source = {
   19:             "type": "jsonl",
   20:             "files": list(files),
   21:             "fields": [list(f) for f in fields],
   22:             "templates": list(templates),
   23:         }
   24:         return self.pipeline
   25: 
   26: 
   27: @dataclass
   28: class TokenizerAPI:
   29:     pipeline: "UnitBuildPipeline"
   30: 
   31:     def sentencepiece(self, *, files: Sequence[str]) -> "UnitBuildPipeline":
   32:         self.pipeline._tokenizer = {
   33:             "type": "sentencepiece",
   34:             "configs": [{"path": p} for p in files],
   35:         }
   36:         return self.pipeline
   37: 
   38:     def hf_json(self, *, files: Sequence[str]) -> "UnitBuildPipeline":
   39:         self.pipeline._tokenizer = {
   40:             "type": "hf_json",
   41:             "configs": [{"path": p} for p in files],
   42:         }
   43:         return self.pipeline
   44: 
   45: 
   46: @dataclass
   47: class OutputAPI:
   48:     pipeline: "UnitBuildPipeline"
   49: 
   50:     def dir(
   51:         self,
   52:         path: str | Sequence[str],
   53:         *,
   54:         samples_per_shard: Sequence[int],
   55:     ) -> "UnitBuildPipeline":
   56:         if isinstance(path, str):
   57:             paths = [path]
   58:         else:
   59:             paths = list(path)
   60:         self.pipeline._output = {
   61:             "type": "dir",
   62:             "paths": paths,
   63:             "samples_per_shard": list(samples_per_shard),
   64:         }
   65:         return self.pipeline
   66: 
   67: 
   68: @dataclass
   69: class UnitBuildPipeline:
   70:     _source: dict | None = None
   71:     _tokenizer: dict | None = None
   72:     _units: dict | None = None
   73:     _output: dict | None = None
   74: 
   75:     source: SourceAPI = field(init=False)
   76:     tokenizer: TokenizerAPI = field(init=False)
   77:     output: OutputAPI = field(init=False)
   78: 
   79:     def __post_init__(self) -> None:
   80:         self.source = SourceAPI(self)
   81:         self.tokenizer = TokenizerAPI(self)
   82:         self.output = OutputAPI(self)
   83: 
   84:     def units(
   85:         self,
   86:         *,
   87:         sizes: Sequence[Sequence[int]],
   88:         build_batch: Sequence[int],
   89:     ) -> "UnitBuildPipeline":
   90:         self._units = {
   91:             "sizes": [list(s) for s in sizes],
   92:             "build_batch": list(build_batch),
   93:         }
   94:         return self
   95: 
   96:     def run(
   97:         self,
   98:         *,
   99:         on_template_unk: str = "prompt",
  100:         on_overflow: str = "prompt",
  101:     ):
  102:         from .runtime import run_unit_build_pipeline
  103:         from mintdim.apis.unit_build.validation import validate_unit_build_pipeline
  104: 
  105:         config = {
  106:             "source": self._source,
  107:             "tokenizer": self._tokenizer,
  108:             "units": self._units,
  109:             "output": self._output,
  110:         }
  111: 
  112:         validate_unit_build_pipeline(config)
  113:         return run_unit_build_pipeline(
  114:             config,
  115:             on_template_unk=on_template_unk,
  116:             on_overflow=on_overflow,
  117:         )
  118: 
  119: 

====================================================================================================
FILE: src\mintdim\apis\unit_build\runtime.py
====================================================================================================
    1: from __future__ import annotations
    2: 
    3: import hashlib
    4: import sys
    5: from bisect import insort
    6: from dataclasses import dataclass, field
    7: from pathlib import Path
    8: from typing import Iterator, Sequence
    9: 
   10: from mintdim.apis.unit_build.ui.logger import log_event
   11: 
   12: from mintdim.apis.unit_build.errors import OutputValidationError, UnitOverflowError
   13: from mintdim.apis.unit_build.config import TokenizerEntry, UnitBuildConfig
   14: from mintdim.apis.unit_build.contracts import TokenizerHandle, TokenizerMetadata
   15: from mintdim.apis.unit_build.adapters.output.directory import prepare_output_dir
   16: from mintdim.apis.unit_build.adapters.output.indices.duplicate import HashCollector, write_duplicate_index
   17: from mintdim.apis.unit_build.adapters.output.histogram import UnitHistogram, write_histogram
   18: from mintdim.apis.unit_build.adapters.output.length_field import write_length_field
   19: from mintdim.apis.unit_build.adapters.output.manifest import write_manifest
   20: from mintdim.apis.unit_build.adapters.output.indices.overflow import OverflowIndexWriter
   21: from mintdim.apis.unit_build.adapters.output.indices.sample import SampleIndexWriter
   22: from mintdim.apis.unit_build.adapters.output.stats import StatsAccumulator, write_stats
   23: from mintdim.apis.unit_build.adapters.output.indices.unknown import UnkIndexWriter
   24: from mintdim.apis.unit_build.adapters.sources.jsonl import iter_jsonl_samples
   25: from mintdim.apis.unit_build.adapters.sources.template import (
   26:     CompiledTemplate,
   27:     LiteralUnkInfo,
   28:     Segment,
   29:     compile_template,
   30:     preflight_template,
   31:     render_segments,
   32: )
   33: from mintdim.apis.unit_build.adapters.tokenizers import load_tokenizer
   34: from mintdim.apis.unit_build.domain.overflow import OverflowAction, OverflowContext, OverflowResolution
   35: from mintdim.apis.unit_build.domain.planner import find_unit
   36: from mintdim.apis.unit_build.adapters.output.shard_writer import (
   37:     ShardRecord,
   38:     ShardWriter,
   39:     assert_no_magic_collision,
   40:     bytes_per_token,
   41:     magic_for_dtype,
   42:     select_token_dtype,
   43: )
   44: 
   45: 
   46: @dataclass(frozen=True)
   47: class _Session:
   48:     files: list[str]
   49:     fields: list[str]
   50:     template: str
   51:     tokenizer_type: str
   52:     tokenizer_entry: TokenizerEntry
   53:     sizes: list[int]
   54:     build_batch: int
   55:     output_dir: str
   56:     samples_per_shard: int
   57:     file_index_for_template: int
   58: 
   59: 
   60: @dataclass(frozen=True)
   61: class _PreparedSession:
   62:     session: _Session
   63:     tokenizer: TokenizerHandle
   64:     compiled: CompiledTemplate
   65:     token_dtype: str
   66:     bytes_per_token: int
   67:     template_unk_findings: list[LiteralUnkInfo]
   68: 
   69: 
   70: @dataclass
   71: class _Policies:
   72:     """Behavior controls plumbed from `.run(on_template_unk=..., on_overflow=...)`."""
   73: 
   74:     template_unk: str  # "prompt" | "abort" | "continue"
   75:     overflow: str  # "prompt" | "abort" | "truncate" | "skip"
   76:     latched_overflow: OverflowResolution | None = None
   77:     latched_extend: bool = False
   78: 
   79: 
   80: @dataclass
   81: class _BatchRecord:
   82:     sample_id: int
   83:     source_file: str
   84:     source_line: int
   85:     segment_tokens: list[list[int]]
   86:     segment_lengths: list[int]
   87:     token_count: int
   88:     unit_size: int
   89:     content_hash: str
   90:     unk_positions: list[int]
   91:     unk_chars: list[str]
   92:     unk_field_indices: list[int]
   93:     empty_fields: list[str]
   94:     shard_path: str = ""
   95: 
   96: 
   97: @dataclass
   98: class _OverflowSkip:
   99:     sample_id: int
  100:     source_file: str
  101:     source_line: int
  102:     token_count: int
  103:     max_unit_size: int
  104:     action: str
  105:     new_unit_size: int | None = None
  106: 
  107: 
  108: def run_unit_build_pipeline(
  109:     raw_or_config,
  110:     *,
  111:     on_template_unk: str = "prompt",
  112:     on_overflow: str = "prompt",
  113: ) -> dict:
  114:     if isinstance(raw_or_config, UnitBuildConfig):
  115:         config = raw_or_config
  116:     else:
  117:         from .validation import validate_unit_build_pipeline
  118: 
  119:         config = validate_unit_build_pipeline(raw_or_config)
  120: 
  121:     sessions = _plan_sessions(config)
  122:     _preflight_output_dirs(sessions)
  123:     prepared = [_prepare_session(session) for session in sessions]
  124:     outputs = []
  125:     for prep in prepared:
  126:         policies = _Policies(template_unk=on_template_unk, overflow=on_overflow)
  127:         if prep.template_unk_findings:
  128:             _resolve_template_unk_policy(prep=prep, policies=policies)
  129:         outputs.append(_run_session(prep, policies=policies))
  130:     return {"outputs": outputs}
  131: 
  132: 
  133: def _plan_sessions(config: UnitBuildConfig) -> list[_Session]:
  134:     if config.is_shared:
  135:         return [
  136:             _Session(
  137:                 files=list(config.source.files),
  138:                 fields=list(config.source.fields[0]),
  139:                 template=config.source.templates[0],
  140:                 tokenizer_type=config.tokenizer.type,
  141:                 tokenizer_entry=config.tokenizer.configs[0],
  142:                 sizes=list(config.units.sizes[0]),
  143:                 build_batch=config.units.build_batch[0],
  144:                 output_dir=config.output.paths[0],
  145:                 samples_per_shard=config.output.samples_per_shard[0],
  146:                 file_index_for_template=0,
  147:             )
  148:         ]
  149: 
  150:     return [
  151:         _Session(
  152:             files=[config.source.files[i]],
  153:             fields=list(_for_file(config.source.fields, i)),
  154:             template=_for_file(config.source.templates, i),
  155:             tokenizer_type=config.tokenizer.type,
  156:             tokenizer_entry=_for_file(config.tokenizer.configs, i),
  157:             sizes=list(_for_file(config.units.sizes, i)),
  158:             build_batch=_for_file(config.units.build_batch, i),
  159:             output_dir=_for_file(config.output.paths, i),
  160:             samples_per_shard=_for_file(config.output.samples_per_shard, i),
  161:             file_index_for_template=i,
  162:         )
  163:         for i in range(config.file_count)
  164:     ]
  165: 
  166: 
  167: def _for_file(values: Sequence, index: int):
  168:     return values[0] if len(values) == 1 else values[index]
  169: 
  170: 
  171: def _preflight_output_dirs(sessions: list[_Session]) -> None:
  172:     for session in sessions:
  173:         path = Path(session.output_dir)
  174:         if path.exists() and any(path.iterdir()):
  175:             raise OutputValidationError(directory=session.output_dir)
  176: 
  177: 
  178: def _prepare_session(session: _Session) -> _PreparedSession:
  179:     tokenizer = load_tokenizer(
  180:         tokenizer_type=session.tokenizer_type,
  181:         entry=session.tokenizer_entry,
  182:     )
  183:     compiled = compile_template(
  184:         file_index=session.file_index_for_template,
  185:         template=session.template,
  186:         declared_fields=session.fields,
  187:         source_files=session.files,
  188:     )
  189:     template_unk_findings = preflight_template(
  190:         compiled=compiled, tokenizer=tokenizer
  191:     )
  192:     token_dtype = select_token_dtype(tokenizer.vocab_size)
  193:     assert_no_magic_collision(
  194:         vocab_size=tokenizer.vocab_size, token_dtype=token_dtype
  195:     )
  196:     return _PreparedSession(
  197:         session=session,
  198:         tokenizer=tokenizer,
  199:         compiled=compiled,
  200:         token_dtype=token_dtype,
  201:         bytes_per_token=bytes_per_token(token_dtype),
  202:         template_unk_findings=template_unk_findings,
  203:     )
  204: 
  205: 
  206: def _resolve_template_unk_policy(*, prep: _PreparedSession, policies: _Policies) -> None:
  207:     """Run P1 decision for this session before the build starts."""
  208:     from mintdim.apis.unit_build.errors import TemplateTokenizerError
  209: 
  210:     file_index = prep.session.file_index_for_template
  211:     if policies.template_unk == "continue":
  212:         return
  213:     if policies.template_unk == "abort":
  214:         raise TemplateTokenizerError(
  215:             file_index=file_index, template=prep.compiled.raw
  216:         )
  217:     if policies.template_unk == "prompt":
  218:         from mintdim.apis.unit_build.ui.prompts import prompt_template_unk, stdin_is_tty
  219: 
  220:         if not stdin_is_tty():
  221:             raise TemplateTokenizerError(
  222:                 file_index=file_index, template=prep.compiled.raw
  223:             )
  224:         proceed = prompt_template_unk(
  225:             file_index=file_index,
  226:             n_findings=len(prep.template_unk_findings),
  227:         )
  228:         if not proceed:
  229:             raise TemplateTokenizerError(
  230:                 file_index=file_index, template=prep.compiled.raw
  231:             )
  232:         return
  233:     raise ValueError(f"unknown on_template_unk value: {policies.template_unk!r}")
  234: 
  235: 
  236: def _resolve_overflow_policy(
  237:     *,
  238:     ctx: OverflowContext,
  239:     policies: _Policies,
  240: ) -> OverflowResolution:
  241:     """Apply the user's overflow policy and return what action to take."""
  242:     if policies.latched_overflow is not None:
  243:         return policies.latched_overflow
  244:     if policies.latched_extend:
  245:         return OverflowResolution(
  246:             action=OverflowAction.EXTEND, new_unit_size=ctx.token_count
  247:         )
  248: 
  249:     if policies.overflow == "abort":
  250:         raise UnitOverflowError(
  251:             sample=f"{ctx.source_file}:{ctx.source_line}",
  252:             token_count=ctx.token_count,
  253:             max_unit_size=ctx.max_unit_size,
  254:         )
  255:     if policies.overflow == "truncate":
  256:         return OverflowResolution(action=OverflowAction.TRUNCATE)
  257:     if policies.overflow == "skip":
  258:         return OverflowResolution(action=OverflowAction.SKIP)
  259: 
  260:     if policies.overflow != "prompt":
  261:         raise ValueError(f"unknown on_overflow value: {policies.overflow!r}")
  262: 
  263:     from mintdim.apis.unit_build.ui.prompts import resolve_overflow_via_prompt, stdin_is_tty
  264: 
  265:     if not stdin_is_tty():
  266:         raise UnitOverflowError(
  267:             sample=f"{ctx.source_file}:{ctx.source_line}",
  268:             token_count=ctx.token_count,
  269:             max_unit_size=ctx.max_unit_size,
  270:         )
  271: 
  272:     resolution, scope = resolve_overflow_via_prompt(ctx=ctx)
  273:     if scope == "all":
  274:         if resolution.action == OverflowAction.EXTEND:
  275:             policies.latched_extend = True
  276:         else:
  277:             policies.latched_overflow = resolution
  278:     return resolution
  279: 
  280: 
  281: def _run_session(prepared: _PreparedSession, *, policies: _Policies) -> dict:
  282:     session = prepared.session
  283:     sizes = list(session.sizes)
  284:     out_dir = prepare_output_dir(session.output_dir)
  285:     n_segments = len(prepared.compiled.segments)
  286: 
  287:     log_event(
  288:         "start",
  289:         output=out_dir,
  290:         files=len(session.files),
  291:         units=",".join(str(size) for size in sizes),
  292:         build_batch=session.build_batch,
  293:     )
  294: 
  295:     shard_writers: dict[int, ShardWriter] = {}
  296:     for unit_size in sizes:
  297:         shard_writers[unit_size] = _make_shard_writer(
  298:             out_dir=out_dir,
  299:             unit_size=unit_size,
  300:             prepared=prepared,
  301:             samples_per_shard=session.samples_per_shard,
  302:             n_segments=n_segments,
  303:         )
  304: 
  305:     sample_index = SampleIndexWriter(out_dir)
  306:     unk_index = UnkIndexWriter(out_dir)
  307:     overflow_index = OverflowIndexWriter(out_dir)
  308:     sample_index.open()
  309:     unk_index.open()
  310:     overflow_index.open()
  311: 
  312:     hashes = HashCollector()
  313:     histogram = UnitHistogram()
  314:     stats = StatsAccumulator(
  315:         token_dtype=prepared.token_dtype,
  316:         bytes_per_token=prepared.bytes_per_token,
  317:     )
  318: 
  319:     sample_id = 0
  320:     batch_index = 0
  321:     success = False
  322:     failure_reason: str | None = None
  323:     try:
  324:         for source_file in session.files:
  325:             for batch in _iter_batches(
  326:                 file_path=source_file,
  327:                 fields=session.fields,
  328:                 batch_size=session.build_batch,
  329:             ):
  330:                 records, skips = _build_batch_records(
  331:                     source_file=source_file,
  332:                     batch=batch,
  333:                     compiled=prepared.compiled,
  334:                     tokenizer=prepared.tokenizer,
  335:                     sample_id_start=sample_id,
  336:                     sizes=sizes,
  337:                     policies=policies,
  338:                     fields=session.fields,
  339:                 )
  340: 
  341:                 _ensure_writers_for_sizes(
  342:                     shard_writers=shard_writers,
  343:                     sizes=sizes,
  344:                     out_dir=out_dir,
  345:                     prepared=prepared,
  346:                     samples_per_shard=session.samples_per_shard,
  347:                     n_segments=n_segments,
  348:                 )
  349:                 _write_records_by_unit(records, shard_writers)
  350:                 sample_index.write_many(_sample_index_rows(records))
  351:                 unk_index.write_many(_unk_index_rows(records))
  352:                 overflow_index.write_many(_overflow_index_rows(skips))
  353:                 hashes.add_many(_hash_rows(records))
  354:                 histogram.add_many([rec.unit_size for rec in records])
  355:                 stats.record_many(
  356:                     token_counts=[rec.token_count for rec in records],
  357:                     unit_sizes=[rec.unit_size for rec in records],
  358:                     unk_counts=[len(rec.unk_positions) for rec in records],
  359:                 )
  360: 
  361:                 sample_id += len(batch)
  362:                 batch_index += 1
  363:                 _log_batch_progress(
  364:                     out_dir=out_dir,
  365:                     batch_index=batch_index,
  366:                     total_samples=sample_id,
  367:                     total_tokens=stats.total_tokens,
  368:                     records=records,
  369:                 )
  370:         success = True
  371:     except BaseException as exc:
  372:         failure_reason = f"{type(exc).__name__}: {exc}"
  373:         log_event(
  374:             "failed",
  375:             output=out_dir,
  376:             batches=batch_index,
  377:             samples=sample_id,
  378:             error=failure_reason,
  379:         )
  380:         raise
  381:     finally:
  382:         for writer in shard_writers.values():
  383:             writer.close(success=success, failure_reason=failure_reason)
  384:         sample_index.close()
  385:         unk_index.close()
  386:         overflow_index.close()
  387: 
  388:     groups = hashes.groups()
  389:     duplicate_groups = sum(1 for group in groups.values() if len(group) > 1)
  390:     duplicate_samples = sum(len(group) for group in groups.values() if len(group) > 1)
  391:     stats.record_duplicates(groups=duplicate_groups, samples=duplicate_samples)
  392: 
  393:     write_duplicate_index(out_dir, groups)
  394:     write_histogram(out_dir, histogram)
  395:     write_stats(out_dir, stats)
  396: 
  397:     magic = magic_for_dtype(prepared.token_dtype)
  398:     write_length_field(
  399:         out_dir=out_dir,
  400:         magic=magic,
  401:         sequence_template=prepared.compiled.sequence_template,
  402:         token_dtype=prepared.token_dtype,
  403:         pad_token_id=prepared.tokenizer.pad_id,
  404:     )
  405: 
  406:     import mintdim
  407: 
  408:     tokenizer_meta = TokenizerMetadata(
  409:         type=session.tokenizer_type,
  410:         path=session.tokenizer_entry.path,
  411:         vocab_size=prepared.tokenizer.vocab_size,
  412:         unk_id=prepared.tokenizer.unk_id,
  413:         pad_id=prepared.tokenizer.pad_id,
  414:     )
  415:     write_manifest(
  416:         out_dir=out_dir,
  417:         mintdim_version=mintdim.__version__,
  418:         tokenizer=tokenizer_meta,
  419:         token_dtype=prepared.token_dtype,
  420:         bytes_per_token=prepared.bytes_per_token,
  421:         magic=magic,
  422:         sequence_template=prepared.compiled.sequence_template,
  423:         sizes=sizes,
  424:         build_batch=session.build_batch,
  425:         samples_per_shard=session.samples_per_shard,
  426:         files=session.files,
  427:         fields=[session.fields],
  428:         templates=[session.template],
  429:     )
  430:     log_event(
  431:         "done",
  432:         output=out_dir,
  433:         batches=batch_index,
  434:         samples=stats.total_samples,
  435:         tokens=stats.total_tokens,
  436:     )
  437: 
  438:     return {
  439:         "output_dir": str(out_dir),
  440:         "total_samples": stats.total_samples,
  441:         "total_tokens": stats.total_tokens,
  442:         "samples_with_unk": stats.samples_with_unk,
  443:         "duplicate_groups": duplicate_groups,
  444:         "sizes": list(sizes),
  445:     }
  446: 
  447: 
  448: def _make_shard_writer(
  449:     *,
  450:     out_dir: Path,
  451:     unit_size: int,
  452:     prepared: _PreparedSession,
  453:     samples_per_shard: int,
  454:     n_segments: int,
  455: ) -> ShardWriter:
  456:     return ShardWriter(
  457:         root=out_dir / f"unit_{unit_size}",
  458:         unit_size=unit_size,
  459:         dtype=prepared.token_dtype,
  460:         pad_id=prepared.tokenizer.pad_id,
  461:         samples_per_shard=samples_per_shard,
  462:         n_segments=n_segments,
  463:     )
  464: 
  465: 
  466: def _ensure_writers_for_sizes(
  467:     *,
  468:     shard_writers: dict[int, ShardWriter],
  469:     sizes: list[int],
  470:     out_dir: Path,
  471:     prepared: _PreparedSession,
  472:     samples_per_shard: int,
  473:     n_segments: int,
  474: ) -> None:
  475:     for unit_size in sizes:
  476:         if unit_size in shard_writers:
  477:             continue
  478:         shard_writers[unit_size] = _make_shard_writer(
  479:             out_dir=out_dir,
  480:             unit_size=unit_size,
  481:             prepared=prepared,
  482:             samples_per_shard=samples_per_shard,
  483:             n_segments=n_segments,
  484:         )
  485: 
  486: 
  487: def _iter_batches(
  488:     *,
  489:     file_path: str,
  490:     fields: list[str],
  491:     batch_size: int,
  492: ) -> Iterator[list[tuple[int, dict]]]:
  493:     buffer: list[tuple[int, dict]] = []
  494:     for line_number, sample in iter_jsonl_samples(file_path, fields):
  495:         buffer.append((line_number, sample))
  496:         if len(buffer) >= batch_size:
  497:             yield buffer
  498:             buffer = []
  499:     if buffer:
  500:         yield buffer
  501: 
  502: 
  503: def _encode_segments(
  504:     *,
  505:     compiled: CompiledTemplate,
  506:     tokenizer: TokenizerHandle,
  507:     sample: dict,
  508: ) -> tuple[list[list[int]], list[str]]:
  509:     """Tokenize each segment of one sample independently.
  510: 
  511:     Returns (segment_token_ids_per_segment, rendered_segment_texts).
  512:     Both lists are parallel to compiled.segments.
  513:     """
  514:     rendered = render_segments(compiled, sample)
  515:     encoded_per_segment = tokenizer.encode_batch(rendered)
  516:     return encoded_per_segment, rendered
  517: 
  518: 
  519: def _build_batch_records(
  520:     *,
  521:     source_file: str,
  522:     batch: list[tuple[int, dict]],
  523:     compiled: CompiledTemplate,
  524:     tokenizer: TokenizerHandle,
  525:     sample_id_start: int,
  526:     sizes: list[int],
  527:     policies: _Policies,
  528:     fields: list[str],
  529: ) -> tuple[list[_BatchRecord], list[_OverflowSkip]]:
  530:     records: list[_BatchRecord] = []
  531:     skips: list[_OverflowSkip] = []
  532: 
  533:     for offset, (line_number, sample) in enumerate(batch):
  534:         sample_id = sample_id_start + offset
  535: 
  536:         segment_tokens, segment_texts = _encode_segments(
  537:             compiled=compiled, tokenizer=tokenizer, sample=sample
  538:         )
  539:         segment_lengths = [len(toks) for toks in segment_tokens]
  540:         token_count = sum(segment_lengths)
  541: 
  542:         empty_fields = [
  543:             seg.name
  544:             for seg, text in zip(compiled.segments, segment_texts)
  545:             if seg.kind == "field" and text == ""
  546:         ]
  547:         content_hash = _content_hash("".join(segment_texts))
  548: 
  549:         fit = find_unit(token_count=token_count, sizes=sizes)
  550:         skip_action: str | None = None
  551:         new_unit_size: int | None = None
  552: 
  553:         if fit is None:
  554:             ctx = OverflowContext(
  555:                 source_file=source_file,
  556:                 source_line=line_number,
  557:                 sample_id=sample_id,
  558:                 token_count=token_count,
  559:                 max_unit_size=sizes[-1],
  560:             )
  561:             resolution = _resolve_overflow_policy(ctx=ctx, policies=policies)
  562: 
  563:             if resolution.action == OverflowAction.SKIP:
  564:                 skip_action = "skip"
  565:                 skips.append(
  566:                     _OverflowSkip(
  567:                         sample_id=sample_id,
  568:                         source_file=source_file,
  569:                         source_line=line_number,
  570:                         token_count=token_count,
  571:                         max_unit_size=sizes[-1],
  572:                         action="skip",
  573:                     )
  574:                 )
  575:                 continue
  576: 
  577:             if resolution.action == OverflowAction.TRUNCATE:
  578:                 skip_action = "truncate"
  579:                 max_size = sizes[-1]
  580:                 segment_tokens, segment_lengths = _truncate_segments(
  581:                     segment_tokens=segment_tokens, max_total=max_size
  582:                 )
  583:                 truncated_token_count = sum(segment_lengths)
  584:                 skips.append(
  585:                     _OverflowSkip(
  586:                         sample_id=sample_id,
  587:                         source_file=source_file,
  588:                         source_line=line_number,
  589:                         token_count=token_count,
  590:                         max_unit_size=max_size,
  591:                         action="truncate",
  592:                     )
  593:                 )
  594:                 token_count = truncated_token_count
  595:                 fit = max_size
  596:             elif resolution.action == OverflowAction.EXTEND:
  597:                 new_unit_size = resolution.new_unit_size
  598:                 assert new_unit_size is not None
  599:                 if new_unit_size not in sizes:
  600:                     insort(sizes, new_unit_size)
  601:                 skips.append(
  602:                     _OverflowSkip(
  603:                         sample_id=sample_id,
  604:                         source_file=source_file,
  605:                         source_line=line_number,
  606:                         token_count=token_count,
  607:                         max_unit_size=ctx.max_unit_size,
  608:                         action="extend",
  609:                         new_unit_size=new_unit_size,
  610:                     )
  611:                 )
  612:                 fit = new_unit_size
  613:             else:  # pragma: no cover â€” exhaustive enum
  614:                 raise ValueError(f"unhandled overflow action: {resolution.action}")
  615: 
  616:         unit_size = fit
  617:         payload_tokens = [tid for seg in segment_tokens for tid in seg]
  618: 
  619:         unk_positions: list[int] = []
  620:         unk_field_indices: list[int] = []
  621:         unk_chars: list[str] = []
  622:         unk_id = tokenizer.unk_id
  623:         cursor = 0
  624:         for seg_index, seg_toks in enumerate(segment_tokens):
  625:             seg_text = segment_texts[seg_index]
  626:             seg_unk_local_positions = [
  627:                 i for i, t in enumerate(seg_toks) if t == unk_id
  628:             ]
  629:             if seg_unk_local_positions:
  630:                 _ids, surfaces = tokenizer.encode_with_offsets(seg_text)
  631:                 for local in seg_unk_local_positions:
  632:                     unk_positions.append(cursor + local)
  633:                     unk_field_indices.append(seg_index)
  634:                     if local < len(surfaces):
  635:                         unk_chars.append(surfaces[local])
  636:                     else:
  637:                         unk_chars.append("")
  638:             cursor += len(seg_toks)
  639: 
  640:         records.append(
  641:             _BatchRecord(
  642:                 sample_id=sample_id,
  643:                 source_file=source_file,
  644:                 source_line=line_number,
  645:                 segment_tokens=segment_tokens,
  646:                 segment_lengths=segment_lengths,
  647:                 token_count=token_count,
  648:                 unit_size=unit_size,
  649:                 content_hash=content_hash,
  650:                 unk_positions=unk_positions,
  651:                 unk_chars=unk_chars,
  652:                 unk_field_indices=unk_field_indices,
  653:                 empty_fields=empty_fields,
  654:             )
  655:         )
  656: 
  657:     return records, skips
  658: 
  659: 
  660: def _truncate_segments(
  661:     *, segment_tokens: list[list[int]], max_total: int
  662: ) -> tuple[list[list[int]], list[int]]:
  663:     """Trim trailing segments so that sum(lengths) == max_total.
  664: 
  665:     Walks segments in order, keeping each full segment until the remaining
  666:     budget can't absorb it; the segment that hits the limit is sliced.
  667:     """
  668:     out: list[list[int]] = []
  669:     budget = max_total
  670:     for toks in segment_tokens:
  671:         if budget <= 0:
  672:             out.append([])
  673:             continue
  674:         if len(toks) <= budget:
  675:             out.append(list(toks))
  676:             budget -= len(toks)
  677:         else:
  678:             out.append(list(toks[:budget]))
  679:             budget = 0
  680:     lengths = [len(toks) for toks in out]
  681:     return out, lengths
  682: 
  683: 
  684: def _write_records_by_unit(
  685:     records: list[_BatchRecord],
  686:     shard_writers: dict[int, ShardWriter],
  687: ) -> None:
  688:     grouped: dict[int, list[_BatchRecord]] = {}
  689:     for record in records:
  690:         grouped.setdefault(record.unit_size, []).append(record)
  691: 
  692:     for unit_size, unit_records in grouped.items():
  693:         shard_records = [
  694:             ShardRecord(
  695:                 segment_lengths=list(rec.segment_lengths),
  696:                 payload_tokens=[t for seg in rec.segment_tokens for t in seg],
  697:             )
  698:             for rec in unit_records
  699:         ]
  700:         locations = shard_writers[unit_size].write_batch(shard_records)
  701:         for record, (shard_index, _position) in zip(unit_records, locations):
  702:             record.shard_path = f"unit_{unit_size}/shard_{shard_index:06d}.bin"
  703: 
  704: 
  705: def _sample_index_rows(records: list[_BatchRecord]) -> list[dict]:
  706:     return [
  707:         {
  708:             "sample_id": record.sample_id,
  709:             "hash": record.content_hash,
  710:             "token_count": record.token_count,
  711:             "unit_size": record.unit_size,
  712:             "source_file": record.source_file,
  713:             "source_line": record.source_line,
  714:             "shard_path": record.shard_path,
  715:             "empty_fields": record.empty_fields,
  716:         }
  717:         for record in records
  718:     ]
  719: 
  720: 
  721: def _unk_index_rows(records: list[_BatchRecord]) -> list[dict]:
  722:     return [
  723:         {
  724:             "sample_id": record.sample_id,
  725:             "hash": record.content_hash,
  726:             "source_file": record.source_file,
  727:             "source_line": record.source_line,
  728:             "unk_count": len(record.unk_positions),
  729:             "unk_positions": record.unk_positions,
  730:             "unk_field_indices": record.unk_field_indices,
  731:             "unk_chars": record.unk_chars,
  732:         }
  733:         for record in records
  734:         if record.unk_positions
  735:     ]
  736: 
  737: 
  738: def _overflow_index_rows(skips: list[_OverflowSkip]) -> list[dict]:
  739:     rows: list[dict] = []
  740:     for skip in skips:
  741:         row: dict = {
  742:             "sample_id": skip.sample_id,
  743:             "source_file": skip.source_file,
  744:             "source_line": skip.source_line,
  745:             "token_count": skip.token_count,
  746:             "max_unit_size": skip.max_unit_size,
  747:             "action": skip.action,
  748:         }
  749:         if skip.new_unit_size is not None:
  750:             row["new_unit_size"] = skip.new_unit_size
  751:         rows.append(row)
  752:     return rows
  753: 
  754: 
  755: def _hash_rows(records: list[_BatchRecord]) -> list[dict]:
  756:     return [
  757:         {
  758:             "sample_id": record.sample_id,
  759:             "hash": record.content_hash,
  760:             "source_file": record.source_file,
  761:             "source_line": record.source_line,
  762:         }
  763:         for record in records
  764:     ]
  765: 
  766: 
  767: def _content_hash(rendered: str) -> str:
  768:     return hashlib.sha256(rendered.encode("utf-8")).hexdigest()
  769: 
  770: 
  771: def _log_batch_progress(
  772:     *,
  773:     out_dir: Path,
  774:     batch_index: int,
  775:     total_samples: int,
  776:     total_tokens: int,
  777:     records: list[_BatchRecord],
  778: ) -> None:
  779:     log_event(
  780:         "progress",
  781:         batch=batch_index,
  782:         batch_samples=len(records),
  783:         total_samples=total_samples,
  784:         batch_tokens=sum(record.token_count for record in records),
  785:         total_tokens=total_tokens,
  786:     )
  787: 
  788: 

====================================================================================================
FILE: src\mintdim\pipeline.py
====================================================================================================
    1: from .apis.registry import create_pipeline
    2: 
    3: def pipeline(name: str):
    4:     return create_pipeline(name)
    5: 
    6: 

====================================================================================================
FILE: src\mintdim\apis\registry.py
====================================================================================================
    1: from .unit_build.api import UnitBuildPipeline
    2: 
    3: 
    4: def create_pipeline(name: str):
    5:     if name == "unit-build":
    6:         return UnitBuildPipeline()
    7:     raise ValueError(f"Unknown pipeline: {name}")
    8: 
    9: 

====================================================================================================
FILE: pyproject.toml
====================================================================================================
    1: [build-system]
    2: requires = ["hatchling"]
    3: build-backend = "hatchling.build"
    4: 
    5: [project]
    6: name = "mintdim"
    7: version = "0.1.13"
    8: description = "Token unit pipeline for building low-padding binary token shards from JSONL datasets."
    9: readme = "README.md"
   10: requires-python = ">=3.12"
   11: license = { file = "LICENSE" }
   12: authors = [
   13:     { name = "MintDim contributors" },
   14: ]
   15: keywords = ["tokenizer", "dataset", "shards", "jsonl", "training"]
   16: classifiers = [
   17:     "Development Status :: 3 - Alpha",
   18:     "Intended Audience :: Developers",
   19:     "Intended Audience :: Science/Research",
   20:     "License :: OSI Approved :: MIT License",
   21:     "Programming Language :: Python :: 3",
   22:     "Programming Language :: Python :: 3.12",
   23:     "Topic :: Scientific/Engineering :: Artificial Intelligence",
   24: ]
   25: dependencies = [
   26:     "numpy>=1.24",
   27:     "sentencepiece>=0.1.99",
   28:     "tokenizers>=0.15",
   29: ]
   30: 
   31: [project.optional-dependencies]
   32: dev = [
   33:     "pytest>=7.4",
   34:     "pytest-cov>=4.1",
   35: ]
   36: 
   37: [project.scripts]
   38: mintdim = "mintdim.cli.main:main"
   39: 
   40: [tool.hatch.build.targets.wheel]
   41: packages = ["src/mintdim"]
   42: 
   43: [tool.pytest.ini_options]
   44: testpaths = ["tests"]
   45: addopts = "-ra --strict-markers --import-mode=importlib"

====================================================================================================
SEARCH: CLI/logger references
====================================================================================================
src\mintdim\apis\unit_build\runtime.py:10: from mintdim.apis.unit_build.ui.logger import log_event
src\mintdim\apis\unit_build\runtime.py:287:     log_event(
src\mintdim\apis\unit_build\runtime.py:373:         log_event(
src\mintdim\apis\unit_build\runtime.py:430:     log_event(
src\mintdim\apis\unit_build\runtime.py:779:     log_event(
src\mintdim\apis\unit_build\adapters\output\shard_writer.py:8: from mintdim.apis.unit_build.ui.logger import log_event
src\mintdim\apis\unit_build\adapters\output\shard_writer.py:195:         log_event("shard", **fields)
src\mintdim\apis\unit_build\adapters\output\shard_writer.py:198:             log_event(
src\mintdim\apis\unit_build\ui\logger.py:58: def main(argv: list[str] | None = None) -> int:
src\mintdim\apis\unit_build\ui\logger.py:78: def log_event(event: str, **fields: Any) -> None:
src\mintdim\apis\unit_build\ui\logger.py:124: def log_result(result: Mapping[str, Any]) -> None:
src\mintdim\apis\unit_build\ui\prompts.py:6: from mintdim.apis.unit_build.ui.logger import (
src\mintdim\cli\main.py:5: from mintdim.cli.commands import unit_build
src\mintdim\cli\main.py:8: def main(argv: list[str] | None = None) -> int:
src\mintdim\cli\main.py:18:         return unit_build.main(args)
src\mintdim\cli\commands\unit_build.py:4: from mintdim.apis.unit_build.ui.logger import main as ui_main
src\mintdim\cli\commands\unit_build.py:11: def main(argv: list[str] | None = None) -> int:
