tools/extend_char_vocab.py

"""Create a derived char vocab with additional path characters."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Append missing characters to an AniFileBERT char vocab JSON."
    )
    parser.add_argument("--input", required=True, help="Base vocab.char.json path")
    parser.add_argument("--output", required=True, help="Derived vocab output path")
    parser.add_argument(
        "--chars",
        default="/\\",
        help="Characters to ensure in the vocab. Default adds slash and backslash.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    input_path = Path(args.input)
    output_path = Path(args.output)
    vocab = json.loads(input_path.read_text(encoding="utf-8"))
    if not isinstance(vocab, dict):
        raise TypeError(f"Expected object vocab JSON: {input_path}")

    next_id = max(int(value) for value in vocab.values()) + 1
    added: list[tuple[str, int]] = []
    for char in args.chars:
        if char not in vocab:
            vocab[char] = next_id
            added.append((char, next_id))
            next_id += 1

    ordered = dict(sorted(vocab.items(), key=lambda item: int(item[1])))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps(ordered, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    print(
        json.dumps(
            {
                "input": str(input_path),
                "output": str(output_path),
                "base_size": len(vocab) - len(added),
                "output_size": len(vocab),
                "added": [{"char": char, "id": idx} for char, idx in added],
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()