Creates a new copy of a gguf file with a different jinja chat template.

update_gguf_jinja.py

#!/usr/bin/env python3
"""
update_gguf_jinja.py — Replace a metadata string key in a GGUF model file.

By default replaces tokenizer.chat_template with the contents of a .jinja
file.  Writes a new output file; the source is never modified.

Usage examples:
    python update_gguf_jinja.py model.gguf template.jinja
    python update_gguf_jinja.py model.gguf template.jinja -o patched.gguf
    python update_gguf_jinja.py model.gguf template.jinja --key tokenizer.chat_template
    python update_gguf_jinja.py model.gguf --show-template

Requirements:
    pip install gguf
    pip install rich           # optional, nicer progress bar
    pip install argcomplete    # optional, enables tab completion
"""

from __future__ import annotations

import argparse
import sys
import threading
import time
from pathlib import Path

try:
    from rich.progress import (
        Progress, BarColumn, FileSizeColumn, TotalFileSizeColumn,
        TransferSpeedColumn, TimeRemainingColumn, TextColumn,
    )
    _HAS_RICH = True
except ImportError:
    _HAS_RICH = False

# ── optional tab-completion ───────────────────────────────────────────────────
try:
    import argcomplete
    _HAS_ARGCOMPLETE = True
except ImportError:
    _HAS_ARGCOMPLETE = False

# ── required dependency ───────────────────────────────────────────────────────
try:
    from gguf import GGUFReader, GGUFWriter, GGUFValueType
except ImportError:
    print("Error: 'gguf' package is required.  Install with:\n    pip install gguf",
          file=sys.stderr)
    sys.exit(1)


# ─────────────────────────────────────────────────────────────────────────────
# Progress bar (file-size based, runs in background thread)
# ─────────────────────────────────────────────────────────────────────────────

class _ProgressMonitor:
    """Polls output file size and updates a progress display every 0.33 s."""

    BAR_WIDTH = 42

    def __init__(self, output_path: Path, total_bytes: int) -> None:
        self._path = output_path
        self._total = total_bytes
        self._stop = threading.Event()

        if _HAS_RICH:
            self._progress = Progress(
                TextColumn("[bold cyan]{task.description}"),
                BarColumn(bar_width=None),
                FileSizeColumn(),
                TextColumn("/"),
                TotalFileSizeColumn(),
                TransferSpeedColumn(),
                TimeRemainingColumn(),
            )
            self._task = self._progress.add_task("Writing tensors", total=total_bytes)
        else:
            self._progress = None
            self._task = None

        self._thread = threading.Thread(target=self._run, daemon=True)

    def start(self) -> None:
        if self._progress:
            self._progress.start()
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        self._thread.join()
        if self._progress:
            self._progress.update(self._task, completed=self._total)
            self._progress.stop()
        else:
            self._draw(self._total)
            print(file=sys.stderr)

    def _run(self) -> None:
        while not self._stop.is_set():
            try:
                current = self._path.stat().st_size
            except FileNotFoundError:
                current = 0
            if self._progress:
                self._progress.update(self._task, completed=current)
            else:
                self._draw(current)
            time.sleep(0.33)

    def _draw(self, current: int) -> None:
        """Fallback ASCII bar used when Rich is not installed."""
        pct = min(current / self._total, 1.0) if self._total > 0 else 1.0
        filled = int(self.BAR_WIDTH * pct)
        bar = "█" * filled + "░" * (self.BAR_WIDTH - filled)
        written_gb = current / 1024 ** 3
        total_gb = self._total / 1024 ** 3
        print(
            f"\r  [{bar}] {written_gb:.2f} / {total_gb:.2f} GB  ({pct:.1%})",
            end="",
            flush=True,
            file=sys.stderr,
        )


# ─────────────────────────────────────────────────────────────────────────────
# Field copy helpers
# ─────────────────────────────────────────────────────────────────────────────

def _copy_field(writer: GGUFWriter, field) -> None:
    """Copy one metadata field from a GGUFReader field to a GGUFWriter."""
    key = field.name
    vtype = field.types[0]

    if vtype == GGUFValueType.STRING:
        writer.add_string(key, str(bytes(field.parts[-1]), encoding="utf-8"))

    elif vtype == GGUFValueType.BOOL:
        writer.add_bool(key, bool(field.parts[-1][0]))

    elif vtype == GGUFValueType.UINT8:
        writer.add_uint8(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.INT8:
        writer.add_int8(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.UINT16:
        writer.add_uint16(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.INT16:
        writer.add_int16(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.UINT32:
        writer.add_uint32(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.INT32:
        writer.add_int32(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.UINT64:
        writer.add_uint64(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.INT64:
        writer.add_int64(key, int(field.parts[-1][0]))
    elif vtype == GGUFValueType.FLOAT32:
        writer.add_float32(key, float(field.parts[-1][0]))
    elif vtype == GGUFValueType.FLOAT64:
        writer.add_float64(key, float(field.parts[-1][0]))

    elif vtype == GGUFValueType.ARRAY:
        elem_type = field.types[1] if len(field.types) > 1 else None

        if elem_type == GGUFValueType.STRING:
            # field.data holds part-indices pointing at each string's byte data
            values = [str(bytes(field.parts[i]), encoding="utf-8") for i in field.data]
            writer.add_array(key, values)

        elif elem_type == GGUFValueType.BOOL:
            writer.add_array(key, [bool(field.parts[i][0]) for i in field.data])

        else:
            # Numeric arrays: GGUFReader stores each element as its own
            # 1-element part; field.data has one index per element — same
            # layout as string/bool arrays.  Collect element [0] from each.
            if field.data:
                writer.add_array(key, [field.parts[i][0].item() for i in field.data])
            else:
                writer.add_array(key, field.parts[-1].tolist())

    else:
        print(f"  Warning: skipping field '{key}' — unsupported type {vtype}",
              file=sys.stderr)


def _get_arch(reader: GGUFReader) -> str:
    f = reader.fields.get("general.architecture")
    if f is None:
        return "unknown"
    return str(bytes(f.parts[-1]), encoding="utf-8")


# ─────────────────────────────────────────────────────────────────────────────
# Core operations
# ─────────────────────────────────────────────────────────────────────────────

def show_template(gguf_path: Path, key: str) -> None:
    reader = GGUFReader(str(gguf_path), "r")
    field = reader.fields.get(key)
    if field is None:
        print(f"Key '{key}' not found in {gguf_path.name}.", file=sys.stderr)
        sys.exit(1)
    vtype = field.types[0]
    if vtype != GGUFValueType.STRING:
        print(f"Key '{key}' is type {vtype.name}, not STRING.", file=sys.stderr)
        sys.exit(1)
    print(str(bytes(field.parts[-1]), encoding="utf-8"))


def patch(
    input_path: Path,
    template_path: Path,
    output_path: Path,
    key: str,
    dry_run: bool,
) -> None:
    new_value = template_path.read_text(encoding="utf-8")

    print(f"  Source : {input_path}", file=sys.stderr)
    print(f"  Template: {template_path}  ({len(new_value):,} chars)", file=sys.stderr)
    print(f"  Key     : {key}", file=sys.stderr)
    print(f"  Output  : {output_path}", file=sys.stderr)

    if dry_run:
        print("\nDry run — nothing written.", file=sys.stderr)
        return

    if output_path.exists():
        print(f"\nError: output file already exists: {output_path}", file=sys.stderr)
        print("Delete it first, or choose a different --output path.", file=sys.stderr)
        sys.exit(1)

    print("\nReading metadata …", file=sys.stderr)
    reader = GGUFReader(str(input_path), "r")
    arch = _get_arch(reader)

    if key not in reader.fields:
        print(f"\nWarning: key '{key}' not found in source; it will be added.",
              file=sys.stderr)

    # Estimate output size ≈ input size + delta of template string
    source_size = input_path.stat().st_size
    old_len = 0
    if key in reader.fields:
        f = reader.fields[key]
        old_len = len(bytes(f.parts[-1]))
    size_estimate = source_size - old_len + len(new_value.encode("utf-8"))

    print(f"  Architecture: {arch}", file=sys.stderr)
    print(f"  Tensors     : {len(reader.tensors):,}", file=sys.stderr)
    print(f"  Source size : {source_size / 1024**3:.2f} GB", file=sys.stderr)
    print(f"  Est. output : {size_estimate / 1024**3:.2f} GB", file=sys.stderr)

    print("\nBuilding writer …", file=sys.stderr)
    writer = GGUFWriter(str(output_path), arch)

    # ── KV metadata ──────────────────────────────────────────────────────────
    replaced = False
    for field_key, field in reader.fields.items():
        if field_key == key:
            writer.add_string(key, new_value)
            replaced = True
        elif field_key == "general.architecture":
            pass  # already set by GGUFWriter constructor
        else:
            _copy_field(writer, field)

    if not replaced:
        writer.add_string(key, new_value)

    # ── Tensor metadata (shapes / types) ─────────────────────────────────────
    print("Registering tensors …", file=sys.stderr)
    for tensor in reader.tensors:
        writer.add_tensor(tensor.name, tensor.data, raw_dtype=tensor.tensor_type)

    # ── Write header + KV ────────────────────────────────────────────────────
    print("Writing header and metadata …", file=sys.stderr)
    writer.write_header_to_file()
    writer.write_kv_data_to_file()

    # ── Write tensor data (the slow part) ────────────────────────────────────
    print("Writing tensor data …", file=sys.stderr)
    monitor = _ProgressMonitor(output_path, size_estimate)
    monitor.start()
    try:
        writer.write_tensors_to_file()
    finally:
        monitor.stop()

    writer.close()

    actual_size = output_path.stat().st_size
    print(f"\nDone.  Output: {output_path}  ({actual_size / 1024**3:.2f} GB)",
          file=sys.stderr)


# ─────────────────────────────────────────────────────────────────────────────
# CLI
# ─────────────────────────────────────────────────────────────────────────────

def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="update_gguf_jinja",
        description="Replace a metadata string key (e.g. chat template) in a GGUF file.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
examples:
  # Replace the chat template and auto-name the output:
  python update_gguf_jinja.py model.gguf Gemma-4-31b-it.v3.jinja

  # Specify the output path explicitly:
  python update_gguf_jinja.py model.gguf template.jinja -o model_v3.gguf

  # Print the current embedded template and exit:
  python update_gguf_jinja.py model.gguf --show-template

  # Preview what would happen without writing anything:
  python update_gguf_jinja.py model.gguf template.jinja --dry-run

tab completion (bash/zsh):
  pip install argcomplete
  eval "$(register-python-argcomplete update_gguf_jinja.py)"
""",
    )

    arg_input = p.add_argument(
        "input",
        metavar="INPUT.gguf",
        type=Path,
        help="Source GGUF model file.",
    )
    arg_template = p.add_argument(
        "template",
        metavar="TEMPLATE.jinja",
        type=Path,
        nargs="?",
        help="Jinja template file to embed.  Required unless --show-template is used.",
    )
    arg_output = p.add_argument(
        "-o", "--output",
        metavar="OUTPUT.gguf",
        type=Path,
        default=None,
        help="Destination GGUF file.  Defaults to <input>_updated.gguf beside the source.",
    )

    if _HAS_ARGCOMPLETE:
        arg_input.completer = argcomplete.completers.FilesCompleter(["*.gguf"])
        arg_template.completer = argcomplete.completers.FilesCompleter(["*.jinja", "*.j2"])
        arg_output.completer = argcomplete.completers.FilesCompleter(["*.gguf"])
    p.add_argument(
        "-k", "--key",
        metavar="KEY",
        default="tokenizer.chat_template",
        help="GGUF metadata key to replace (default: tokenizer.chat_template).",
    )
    p.add_argument(
        "--show-template",
        action="store_true",
        help="Print the current value of --key from INPUT.gguf and exit.",
    )
    p.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be done without writing any output.",
    )

    if _HAS_ARGCOMPLETE:
        argcomplete.autocomplete(p)

    return p


def main() -> None:
    parser = _build_parser()
    args = parser.parse_args()

    # ── Validate input ────────────────────────────────────────────────────────
    if not args.input.exists():
        parser.error(f"Input file not found: {args.input}")
    if not args.input.suffix.lower() == ".gguf":
        print(f"Warning: '{args.input}' does not have a .gguf extension.",
              file=sys.stderr)

    if args.show_template:
        show_template(args.input, args.key)
        return

    if args.template is None:
        parser.error("TEMPLATE.jinja is required (or use --show-template).")
    if not args.template.exists():
        parser.error(f"Template file not found: {args.template}")

    # ── Resolve output path ───────────────────────────────────────────────────
    output = args.output
    if output is None:
        stem = args.input.stem
        output = args.input.with_name(f"{stem}_updated.gguf")

    # ── Run ───────────────────────────────────────────────────────────────────
    patch(
        input_path=args.input,
        template_path=args.template,
        output_path=output,
        key=args.key,
        dry_run=args.dry_run,
    )


if __name__ == "__main__":
    main()