cute-tokenizer 0.1.0__py3-none-any.whl

cute_tokenizer/__init__.py

@@ -0,0 +1,35 @@
+"""CUTE — Compact Unicode Token Encoding.
+
+Public API:
+    build_cute      — train a CUTE tokenizer from a corpus directory.
+    CUTEConfig      — all knobs for the build pipeline.
+    CUTETokenizerFast — HuggingFace-compatible inference wrapper.
+    PUAMapping      — word ↔ PUA character mapping.
+"""
+
+from __future__ import annotations
+
+import os as _os
+
+# Silence the "None of PyTorch, TensorFlow, or Flax have been found" warning
+# that `transformers` emits at import. CUTE only needs the tokenizer layer, not
+# the model layer, so this warning is irrelevant. Must be set BEFORE the first
+# `transformers` import — putting it here covers both library and CLI paths.
+_os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+_os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
+
+from ._version import __version__
+from .config import CUTEConfig
+from .pua import PUAMapping
+from .tokenizer import CUTETokenizerFast
+from .trainer import build_cute, load_mapping, save_mapping
+
+__all__ = [
+    "CUTEConfig",
+    "CUTETokenizerFast",
+    "PUAMapping",
+    "__version__",
+    "build_cute",
+    "load_mapping",
+    "save_mapping",
+]

cute_tokenizer/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

cute_tokenizer/cli.py

@@ -0,0 +1,130 @@
+"""Command-line interface for CUTE.
+
+cute build --corpus ./corpus --output ./output [--config configs/default.toml]
+cute roundtrip-check --tokenizer ./output --corpus ./holdout
+cute info --tokenizer ./output
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+# Silence the noisy "None of PyTorch, TensorFlow >= 2.0, or Flax have been found"
+# warning from `transformers`. We only use the tokenizer, not the model layer,
+# so this warning is irrelevant. Set BEFORE any transformers import below.
+os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
+os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
+
+from ._version import __version__
+from .config import CUTEConfig
+from .corpus import iter_corpus_files
+from .manifest import BuildManifest
+from .tokenizer import CUTETokenizerFast
+from .trainer import build_cute
+
+
+def _cmd_build(args: argparse.Namespace) -> int:
+    config = _load_config(Path(args.config)) if args.config else CUTEConfig()
+    manifest_path = build_cute(
+        corpus_dir=Path(args.corpus),
+        output_dir=Path(args.output),
+        config=config,
+    )
+    print(f"Build complete. Manifest: {manifest_path}")
+    return 0
+
+
+def _cmd_roundtrip_check(args: argparse.Namespace) -> int:
+    tok_dir = Path(args.tokenizer)
+    tok = CUTETokenizerFast(
+        tokenizer_file=tok_dir / "tokenizer.json",
+        cute_mapping_file=tok_dir / "cute_mapping.json",
+    )
+
+    corpus_dir = Path(args.corpus)
+    files_checked = files_failed = 0
+    for path in iter_corpus_files(
+        corpus_dir,
+        extensions=(".py", ".js", ".ts", ".java", ".c", ".cpp", ".rs", ".go", ".rb", ".php"),
+    ):
+        try:
+            text = path.read_text(encoding="utf-8")
+        except (UnicodeDecodeError, OSError):
+            continue
+        files_checked += 1
+        ids = tok(text, add_special_tokens=False).input_ids
+        decoded = tok.decode(ids, skip_special_tokens=True)
+        if decoded != text:
+            files_failed += 1
+            print(f"FAIL {path.relative_to(corpus_dir)}")
+            if args.verbose:
+                _show_diff(text, decoded)
+        if files_checked >= args.max_files:
+            break
+
+    print(f"Round-trip check: {files_checked - files_failed}/{files_checked} OK")
+    return 0 if files_failed == 0 else 1
+
+
+def _cmd_info(args: argparse.Namespace) -> int:
+    tok_dir = Path(args.tokenizer)
+    manifest = BuildManifest.read(tok_dir / "build_manifest.json")
+    print(json.dumps(manifest.to_dict(), indent=2))
+    return 0
+
+
+def _show_diff(expected: str, got: str) -> None:
+    """Print a brief diff for round-trip failures."""
+    for i, (e, g) in enumerate(zip(expected, got, strict=False)):
+        if e != g:
+            ctx = max(0, i - 20)
+            print(f"  first diff at offset {i}:")
+            print(f"    expected: ...{expected[ctx : i + 20]!r}")
+            print(f"    got     : ...{got[ctx : i + 20]!r}")
+            return
+    if len(expected) != len(got):
+        print(f"  length differs: expected={len(expected)}, got={len(got)}")
+
+
+def _load_config(path: Path) -> CUTEConfig:
+    """Load a config from TOML. Lazy import — `tomllib` is stdlib in 3.11+."""
+    if sys.version_info >= (3, 11):
+        import tomllib
+    else:
+        import tomli as tomllib  # type: ignore[no-redef]
+    data = tomllib.loads(path.read_text(encoding="utf-8"))
+    return CUTEConfig(**data)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="cute", description="CUTE tokenizer builder")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    p_build = sub.add_parser("build", help="Train a CUTE tokenizer from a corpus")
+    p_build.add_argument("--corpus", required=True, help="Corpus directory")
+    p_build.add_argument("--output", required=True, help="Output directory")
+    p_build.add_argument("--config", help="Optional TOML config path")
+    p_build.set_defaults(func=_cmd_build)
+
+    p_rt = sub.add_parser("roundtrip-check", help="Verify byte-equal round-trip")
+    p_rt.add_argument("--tokenizer", required=True, help="Trained tokenizer dir")
+    p_rt.add_argument("--corpus", required=True, help="Held-out corpus to check")
+    p_rt.add_argument("--max-files", type=int, default=10_000)
+    p_rt.add_argument("--verbose", action="store_true")
+    p_rt.set_defaults(func=_cmd_roundtrip_check)
+
+    p_info = sub.add_parser("info", help="Print build manifest")
+    p_info.add_argument("--tokenizer", required=True)
+    p_info.set_defaults(func=_cmd_info)
+
+    args = parser.parse_args(argv)
+    return int(args.func(args))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cute_tokenizer/config.py

@@ -0,0 +1,96 @@
+"""Configuration for the CUTE tokenizer build pipeline."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import Any
+
+DEFAULT_SPECIAL_TOKENS: tuple[str, ...] = (
+    "<pad>",
+    "<s>",
+    "</s>",
+    "<unk>",
+    "<|endoftext|>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|file_sep|>",
+    "<|repo_name|>",
+)
+
+DEFAULT_CODE_EXTENSIONS: tuple[str, ...] = (
+    ".py",
+    ".js",
+    ".ts",
+    ".tsx",
+    ".jsx",
+    ".java",
+    ".c",
+    ".cpp",
+    ".h",
+    ".hpp",
+    ".cs",
+    ".rs",
+    ".go",
+    ".rb",
+    ".php",
+    ".swift",
+    ".kt",
+    ".scala",
+    ".sh",
+    ".sql",
+    ".html",
+    ".css",
+    ".scss",
+    ".json",
+    ".yaml",
+    ".yml",
+    ".toml",
+    ".md",
+)
+
+
+@dataclass(frozen=True)
+class CUTEConfig:
+    """All knobs for a CUTE build, in one place.
+
+    Frozen so hashing/comparison is well-defined and the manifest serializer
+    can dump a stable representation.
+    """
+
+    vocab_size: int = 80_000
+    coverage_target: float = 0.90
+    max_token_len: int = 50
+    boost_weight: float = 0.3
+    min_bpe_budget: int = 8_000
+    min_frequency: int = 2
+    seed: int = 42
+    extensions: tuple[str, ...] = DEFAULT_CODE_EXTENSIONS
+    special_tokens: tuple[str, ...] = DEFAULT_SPECIAL_TOKENS
+    workers: int = 0  # 0 means os.cpu_count()
+    shard_size_bytes: int = 64 * 1024 * 1024  # 64 MiB per shard
+    license_allowlist: tuple[str, ...] = (
+        "MIT",
+        "Apache-2.0",
+        "BSD-3-Clause",
+        "BSD-2-Clause",
+        "ISC",
+        "Apache 2.0",
+        "Apache License 2.0",
+    )
+    enable_secret_scrub: bool = True
+    enable_license_filter: bool = False  # off by default; opt-in
+
+    def __post_init__(self) -> None:
+        if not 0.0 < self.coverage_target < 1.0:
+            raise ValueError(f"coverage_target must be in (0,1), got {self.coverage_target}")
+        if self.vocab_size < 1024:
+            raise ValueError(f"vocab_size too small: {self.vocab_size}")
+        if self.max_token_len < 1:
+            raise ValueError(f"max_token_len must be positive: {self.max_token_len}")
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+__all__ = ["DEFAULT_CODE_EXTENSIONS", "DEFAULT_SPECIAL_TOKENS", "CUTEConfig"]

cute_tokenizer/corpus.py

@@ -0,0 +1,305 @@
+"""Corpus pipeline: stream files, dedupe by content hash, scrub secrets, shard.
+
+The output is a sequence of deterministic shards on disk that downstream
+phases (frequency counting, BPE training) can iterate efficiently.
+"""
+
+from __future__ import annotations
+
+import gzip
+import hashlib
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import orjson
+import regex as re
+
+from .pua import find_pua_codepoints
+
+# ---------------------------------------------------------------------------
+# Secret scrubbing
+# ---------------------------------------------------------------------------
+
+# Each pattern is conservative — false positives drop a file, which is fine
+# at corpus scale. False negatives are far more dangerous (secret in vocab).
+SECRET_PATTERNS: tuple[tuple[str, re.Pattern[str]], ...] = (
+    ("aws_access_key", re.compile(r"AKIA[0-9A-Z]{16}")),
+    ("openai_api_key", re.compile(r"sk-(?:proj-)?[A-Za-z0-9_-]{20,}")),
+    ("anthropic_api_key", re.compile(r"sk-ant-[A-Za-z0-9_\-]{50,}")),
+    ("github_pat", re.compile(r"ghp_[A-Za-z0-9]{36}")),
+    ("github_oauth", re.compile(r"gho_[A-Za-z0-9]{36}")),
+    ("github_app", re.compile(r"(ghu|ghs)_[A-Za-z0-9]{36}")),
+    ("google_api", re.compile(r"AIza[0-9A-Za-z_\-]{35}")),
+    ("slack_token", re.compile(r"xox[baprs]-[A-Za-z0-9-]{10,}")),
+    ("private_key_pem", re.compile(r"-----BEGIN (?:RSA |EC |DSA |OPENSSH )?PRIVATE KEY-----")),
+    ("jwt", re.compile(r"eyJ[A-Za-z0-9_\-]{10,}\.eyJ[A-Za-z0-9_\-]{10,}\.[A-Za-z0-9_\-]{10,}")),
+)
+
+
+def has_secret(text: str) -> str | None:
+    """Return the name of the first matching secret pattern, or None."""
+    for name, pat in SECRET_PATTERNS:
+        if pat.search(text):
+            return name
+    return None
+
+
+# ---------------------------------------------------------------------------
+# License filter
+# ---------------------------------------------------------------------------
+
+# SPDX header detector. Matches lines like:
+#   # SPDX-License-Identifier: MIT
+#   // SPDX-License-Identifier: Apache-2.0
+# Operates on the first 4 KiB of each file so we don't scan large blobs.
+_SPDX_REGEX = re.compile(
+    r"SPDX-License-Identifier\s*:\s*([A-Za-z0-9.\-+ ]+)",
+    re.IGNORECASE,
+)
+# Heuristic: explicit "All rights reserved" / "Proprietary" / "Confidential"
+# in the file head. We refuse files matching these unless an SPDX header
+# explicitly grants a permissive license.
+_PROPRIETARY_REGEX = re.compile(
+    r"\b(All Rights Reserved|Proprietary and Confidential|UNLICENSED|License: Proprietary)\b",
+    re.IGNORECASE,
+)
+_LICENSE_HEAD_BYTES = 4096
+
+
+def detect_license(text: str) -> str | None:
+    """Best-effort license detection from a file's head.
+
+    Returns the SPDX identifier if found, otherwise None. Does NOT make a
+    keep/drop decision — that's `is_license_allowed`'s job.
+    """
+    head = text[:_LICENSE_HEAD_BYTES]
+    m = _SPDX_REGEX.search(head)
+    if m:
+        return m.group(1).strip()
+    return None
+
+
+def is_license_allowed(text: str, allowlist: Iterable[str]) -> bool:
+    """Decide whether a file's license header (if any) permits inclusion.
+
+    Logic:
+      1. If an SPDX header is present and matches the allowlist → allow.
+      2. If an SPDX header is present and does NOT match → reject.
+      3. If no SPDX header but a 'proprietary' marker is in the head → reject.
+      4. Otherwise (no headers, no markers) → allow. The corpus owner is
+         responsible for not feeding obviously copyrighted material; the
+         filter is a safety net, not a legal review.
+    """
+    spdx = detect_license(text)
+    if spdx is not None:
+        allow_set = {entry.strip().lower() for entry in allowlist}
+        return spdx.lower() in allow_set
+
+    head = text[:_LICENSE_HEAD_BYTES]
+    return not _PROPRIETARY_REGEX.search(head)
+
+
+# ---------------------------------------------------------------------------
+# Records
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class CorpusRecord:
+    """One file's content + metadata."""
+
+    path: str  # path relative to the corpus root
+    text: str
+    sha256: str
+
+    def to_json(self) -> bytes:
+        return orjson.dumps({"path": self.path, "text": self.text, "sha256": self.sha256})
+
+    @classmethod
+    def from_json(cls, line: bytes) -> CorpusRecord:
+        d = orjson.loads(line)
+        return cls(path=d["path"], text=d["text"], sha256=d["sha256"])
+
+
+@dataclass(frozen=True)
+class IngestStats:
+    """Aggregated stats from one ingest pass."""
+
+    files_seen: int
+    files_kept: int
+    files_dropped_dedup: int
+    files_dropped_secret: int
+    files_dropped_license: int
+    files_dropped_decode: int
+    files_dropped_size: int
+    bytes_kept: int
+    pua_codepoints_in_corpus: frozenset[int]
+
+
+# ---------------------------------------------------------------------------
+# Ingestion
+# ---------------------------------------------------------------------------
+
+
+def _hash_text(text: str) -> str:
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def iter_corpus_files(
+    corpus_dir: Path,
+    extensions: Iterable[str],
+    max_bytes: int = 5_000_000,
+) -> Iterator[Path]:
+    """Yield candidate files under `corpus_dir`, deterministically ordered.
+
+    Sorted by relative path so iteration order is reproducible across runs.
+    """
+    ext_set = {e.lower() for e in extensions}
+    candidates = [p for p in corpus_dir.rglob("*") if p.is_file() and p.suffix.lower() in ext_set]
+    candidates.sort(key=lambda p: str(p.relative_to(corpus_dir)).replace("\\", "/"))
+    for p in candidates:
+        try:
+            if p.stat().st_size > max_bytes:
+                continue
+        except OSError:
+            continue
+        yield p
+
+
+def ingest_corpus(
+    corpus_dir: Path,
+    out_dir: Path,
+    extensions: Iterable[str],
+    shard_size_bytes: int = 64 * 1024 * 1024,
+    enable_secret_scrub: bool = True,
+    enable_license_filter: bool = False,
+    license_allowlist: Iterable[str] = (),
+    max_file_bytes: int = 5_000_000,
+) -> IngestStats:
+    """Read corpus files, dedupe + scrub, write line-delimited gzipped shards.
+
+    Output layout:
+        out_dir/shards/shard_00000.jsonl.gz
+        out_dir/shards/shard_00001.jsonl.gz
+        ...
+
+    Each line of each shard is a CorpusRecord.to_json() blob.
+    """
+    out_dir.mkdir(parents=True, exist_ok=True)
+    shards_dir = out_dir / "shards"
+    shards_dir.mkdir(exist_ok=True)
+
+    seen_hashes: set[str] = set()
+    files_seen = files_kept = 0
+    drop_dedup = drop_secret = drop_license = drop_decode = drop_size = 0
+    bytes_kept = 0
+    pua_codepoints: set[int] = set()
+    license_allowlist_t = tuple(license_allowlist)
+
+    shard_idx = 0
+    shard_path = shards_dir / f"shard_{shard_idx:05d}.jsonl.gz"
+    shard_fh: gzip.GzipFile | None = gzip.open(shard_path, "wb")  # noqa: SIM115 (rolling handle, closed in finally)
+    bytes_in_shard = 0
+
+    try:
+        for path in iter_corpus_files(corpus_dir, extensions, max_bytes=max_file_bytes):
+            files_seen += 1
+            try:
+                text = path.read_text(encoding="utf-8", errors="strict")
+            except (UnicodeDecodeError, OSError):
+                drop_decode += 1
+                continue
+
+            if not text:
+                drop_size += 1
+                continue
+
+            sha = _hash_text(text)
+            if sha in seen_hashes:
+                drop_dedup += 1
+                continue
+
+            if enable_secret_scrub and has_secret(text):
+                drop_secret += 1
+                continue
+
+            if enable_license_filter and not is_license_allowed(text, license_allowlist_t):
+                drop_license += 1
+                continue
+
+            seen_hashes.add(sha)
+            pua_codepoints.update(find_pua_codepoints(text))
+
+            rec = CorpusRecord(
+                path=str(path.relative_to(corpus_dir)).replace("\\", "/"),
+                text=text,
+                sha256=sha,
+            )
+            line = rec.to_json() + b"\n"
+
+            assert shard_fh is not None
+            if bytes_in_shard + len(line) > shard_size_bytes and bytes_in_shard > 0:
+                shard_fh.close()
+                shard_idx += 1
+                shard_path = shards_dir / f"shard_{shard_idx:05d}.jsonl.gz"
+                shard_fh = gzip.open(shard_path, "wb")  # noqa: SIM115
+                bytes_in_shard = 0
+
+            shard_fh.write(line)
+            bytes_in_shard += len(line)
+            bytes_kept += len(text.encode("utf-8"))
+            files_kept += 1
+    finally:
+        if shard_fh is not None:
+            shard_fh.close()
+
+    return IngestStats(
+        files_seen=files_seen,
+        files_kept=files_kept,
+        files_dropped_dedup=drop_dedup,
+        files_dropped_secret=drop_secret,
+        files_dropped_license=drop_license,
+        files_dropped_decode=drop_decode,
+        files_dropped_size=drop_size,
+        bytes_kept=bytes_kept,
+        pua_codepoints_in_corpus=frozenset(pua_codepoints),
+    )
+
+
+def iter_shards(shards_dir: Path) -> Iterator[Path]:
+    """Yield shard paths in deterministic order."""
+    shards = sorted(shards_dir.glob("shard_*.jsonl.gz"))
+    yield from shards
+
+
+def read_shard(shard_path: Path) -> Iterator[CorpusRecord]:
+    """Stream records from one shard."""
+    with gzip.open(shard_path, "rb") as fh:
+        for line in fh:
+            line = line.strip()
+            if not line:
+                continue
+            yield CorpusRecord.from_json(line)
+
+
+def iter_shard_texts(shards_dir: Path) -> Iterator[str]:
+    """Stream text payloads from all shards in order."""
+    for shard in iter_shards(shards_dir):
+        for rec in read_shard(shard):
+            yield rec.text
+
+
+__all__ = [
+    "SECRET_PATTERNS",
+    "CorpusRecord",
+    "IngestStats",
+    "detect_license",
+    "has_secret",
+    "ingest_corpus",
+    "is_license_allowed",
+    "iter_corpus_files",
+    "iter_shard_texts",
+    "iter_shards",
+    "read_shard",
+]

cute_tokenizer/decode.py

@@ -0,0 +1,37 @@
+"""PUA-aware decoding.
+
+After the underlying ByteLevel BPE decoder reconstructs a string, we must
+substitute every PUA character back to its original word. This is a single
+linear scan with a dict lookup per character — O(n).
+"""
+
+from __future__ import annotations
+
+from .pua import PUAMapping, is_pua_char
+
+
+def reverse_pua_substitute(text: str, mapping: PUAMapping) -> str:
+    """Replace every PUA character in `text` with its original mapped word.
+
+    Characters not in the mapping are passed through unchanged. This is
+    safe even if the input contains PUA chars that weren't in the mapping
+    (they survive the round-trip as themselves).
+    """
+    pua_to_word = mapping.pua_to_word
+    if not pua_to_word:
+        return text
+
+    # Fast path: if no PUA chars present, return as-is.
+    if not any(is_pua_char(c) for c in text):
+        return text
+
+    out: list[str] = []
+    for ch in text:
+        if is_pua_char(ch):
+            out.append(pua_to_word.get(ch, ch))
+        else:
+            out.append(ch)
+    return "".join(out)
+
+
+__all__ = ["reverse_pua_substitute"]