ckptkit 0.3.0__py3-none-any.whl

ckptkit/__init__.py

@@ -0,0 +1,112 @@
+"""ckptkit — inspect, convert, diff, and merge model checkpoints."""
+
+from ckptkit._types import (
+    DTYPE_SIZES,
+    CheckpointFormat,
+    CheckpointInfo,
+    CkptError,
+    DiffEntry,
+    DiffResult,
+    DType,
+    FormatError,
+    MergeConfig,
+    MergeError,
+    TensorInfo,
+)
+from ckptkit.convert import (
+    ConversionConfig,
+    ConversionFormat,
+    ConversionResult,
+    FormatConverter,
+    convert_checkpoint,
+)
+from ckptkit.diff import diff, diff_infos, format_diff, format_diff_rich, format_diff_table
+from ckptkit.estimator import (
+    EstimationResult,
+    QuantEstimationResult,
+    TensorEstimate,
+    estimate_quantized_size,
+    estimate_reduction,
+    format_estimation,
+)
+from ckptkit.gguf import GGUFInfo, format_gguf_info, inspect_gguf, parse_gguf
+from ckptkit.inspect import (
+    detect_format,
+    format_params,
+    format_size,
+    inspect,
+    inspect_safetensors,
+)
+from ckptkit.merge import find_lora_pairs, merge_lora_state_dicts
+from ckptkit.metadata import (
+    CheckpointMetadata,
+    MetadataEditor,
+    extract_metadata_from_path,
+    format_metadata_report,
+)
+from ckptkit.stats import CheckpointStats, TensorStats, stats_from_info
+from ckptkit.validate import ValidationIssue, ValidationResult, validate
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "__version__",
+    # Types
+    "CheckpointFormat",
+    "CheckpointInfo",
+    "CkptError",
+    "DType",
+    "DTYPE_SIZES",
+    "DiffEntry",
+    "DiffResult",
+    "FormatError",
+    "MergeConfig",
+    "MergeError",
+    "TensorInfo",
+    # Inspect
+    "detect_format",
+    "inspect",
+    "inspect_safetensors",
+    "format_size",
+    "format_params",
+    # Diff
+    "diff",
+    "diff_infos",
+    "format_diff",
+    "format_diff_rich",
+    "format_diff_table",
+    # Merge
+    "merge_lora_state_dicts",
+    "find_lora_pairs",
+    # Stats
+    "CheckpointStats",
+    "TensorStats",
+    "stats_from_info",
+    # Validate
+    "validate",
+    "ValidationResult",
+    "ValidationIssue",
+    # Estimator
+    "EstimationResult",
+    "QuantEstimationResult",
+    "TensorEstimate",
+    "estimate_reduction",
+    "estimate_quantized_size",
+    "format_estimation",
+    # Convert
+    "ConversionConfig",
+    "ConversionFormat",
+    "ConversionResult",
+    "FormatConverter",
+    "convert_checkpoint",
+    # GGUF
+    "GGUFInfo",
+    "parse_gguf",
+    "inspect_gguf",
+    "format_gguf_info",
+    # Metadata
+    "CheckpointMetadata",
+    "MetadataEditor",
+    "extract_metadata_from_path",
+    "format_metadata_report",
+]

ckptkit/_types.py

@@ -0,0 +1,164 @@
+"""Core types for ckptkit."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class CheckpointFormat(str, Enum):
+    """Supported checkpoint formats."""
+
+    SAFETENSORS = "safetensors"
+    PYTORCH = "pytorch"
+    NUMPY = "numpy"
+    UNKNOWN = "unknown"
+
+
+class DType(str, Enum):
+    """Common tensor data types."""
+
+    F32 = "F32"
+    F16 = "F16"
+    BF16 = "BF16"
+    F64 = "F64"
+    I64 = "I64"
+    I32 = "I32"
+    I16 = "I16"
+    I8 = "I8"
+    U8 = "U8"
+    BOOL = "BOOL"
+    UNKNOWN = "UNKNOWN"
+
+
+# Bytes per element for each dtype
+DTYPE_SIZES: dict[DType, int] = {
+    DType.F32: 4, DType.F16: 2, DType.BF16: 2, DType.F64: 8,
+    DType.I64: 8, DType.I32: 4, DType.I16: 2, DType.I8: 1,
+    DType.U8: 1, DType.BOOL: 1,
+}
+
+
+@dataclass
+class TensorInfo:
+    """Metadata for a single tensor in a checkpoint."""
+
+    name: str
+    shape: list[int]
+    dtype: DType
+    offset_start: int = 0
+    offset_end: int = 0
+
+    @property
+    def numel(self) -> int:
+        """Number of elements."""
+        result = 1
+        for s in self.shape:
+            result *= s
+        return result
+
+    @property
+    def size_bytes(self) -> int:
+        """Size in bytes."""
+        return self.numel * DTYPE_SIZES.get(self.dtype, 0)
+
+    @property
+    def shape_str(self) -> str:
+        return "×".join(str(s) for s in self.shape)
+
+
+@dataclass
+class CheckpointInfo:
+    """Metadata for an entire checkpoint."""
+
+    path: str
+    format: CheckpointFormat
+    file_size: int
+    tensors: list[TensorInfo]
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def n_tensors(self) -> int:
+        return len(self.tensors)
+
+    @property
+    def n_parameters(self) -> int:
+        return sum(t.numel for t in self.tensors)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(t.size_bytes for t in self.tensors)
+
+    def dtype_summary(self) -> dict[str, int]:
+        """Count parameters per dtype."""
+        counts: dict[str, int] = {}
+        for t in self.tensors:
+            key = t.dtype.value
+            counts[key] = counts.get(key, 0) + t.numel
+        return counts
+
+    def layer_groups(self) -> dict[str, list[TensorInfo]]:
+        """Group tensors by layer prefix (e.g., 'model.layers.0')."""
+        groups: dict[str, list[TensorInfo]] = {}
+        for t in self.tensors:
+            parts = t.name.split(".")
+            # Take first 3 parts as group key, or full name if shorter
+            key = ".".join(parts[:3]) if len(parts) > 3 else t.name
+            groups.setdefault(key, []).append(t)
+        return groups
+
+
+@dataclass
+class DiffEntry:
+    """A difference between two checkpoints."""
+
+    tensor_name: str
+    change_type: str  # "added", "removed", "shape_changed", "dtype_changed", "values_changed"
+    details: str = ""
+
+
+@dataclass
+class DiffResult:
+    """Result of comparing two checkpoints."""
+
+    path_a: str
+    path_b: str
+    entries: list[DiffEntry]
+    n_shared: int = 0
+    n_identical: int = 0
+
+    @property
+    def n_changes(self) -> int:
+        return len(self.entries)
+
+    @property
+    def has_changes(self) -> bool:
+        return len(self.entries) > 0
+
+
+@dataclass
+class MergeConfig:
+    """Configuration for LoRA adapter merging."""
+
+    base_path: str
+    adapter_path: str
+    output_path: str
+    alpha: float = 1.0
+    device: str = "cpu"
+
+    def __post_init__(self) -> None:
+        if not (0.0 <= self.alpha <= 2.0):
+            raise ValueError(f"alpha must be between 0.0 and 2.0, got {self.alpha}")
+
+
+class CkptError(Exception):
+    """Base exception for ckptkit."""
+
+
+class FormatError(CkptError):
+    """Unsupported or corrupt file format."""
+
+
+class MergeError(CkptError):
+    """Error during checkpoint merging."""

ckptkit/cli.py

@@ -0,0 +1,163 @@
+"""CLI for ckptkit."""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+try:
+    import click
+    _HAS_CLICK = True
+except ImportError:
+    _HAS_CLICK = False
+
+try:
+    from rich.console import Console
+    from rich.table import Table
+    _console = Console()
+    _HAS_RICH = True
+except ImportError:
+    _HAS_RICH = False
+    _console = None  # type: ignore[assignment]
+
+
+def _build_cli() -> Any:
+    if not _HAS_CLICK:
+        return None
+
+    from ckptkit.diff import diff, format_diff
+    from ckptkit.inspect import format_params, format_size, inspect
+    from ckptkit.stats import format_stats, stats_from_info
+    from ckptkit.validate import validate
+
+    @click.group()
+    @click.version_option(package_name="ckptkit")
+    def cli() -> None:
+        """ckptkit — inspect, convert, diff, and merge model checkpoints."""
+
+    @cli.command()
+    @click.argument("path", type=click.Path(exists=True))
+    @click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
+    def info(path: str, as_json: bool) -> None:
+        """Inspect a checkpoint file."""
+        import json as json_mod
+        ckpt_info = inspect(path)
+
+        if as_json:
+            data = {
+                "path": ckpt_info.path,
+                "format": ckpt_info.format.value,
+                "file_size": ckpt_info.file_size,
+                "n_tensors": ckpt_info.n_tensors,
+                "n_parameters": ckpt_info.n_parameters,
+                "tensors": [
+                    {"name": t.name, "shape": t.shape, "dtype": t.dtype.value}
+                    for t in ckpt_info.tensors
+                ],
+            }
+            click.echo(json_mod.dumps(data, indent=2))
+            return
+
+        if _HAS_RICH and _console is not None:
+            _console.print(f"\n[bold]File:[/bold] {ckpt_info.path}")
+            _console.print(f"[bold]Format:[/bold] {ckpt_info.format.value}")
+            _console.print(f"[bold]Size:[/bold] {format_size(ckpt_info.file_size)}")
+            _console.print(f"[bold]Parameters:[/bold] {ckpt_info.n_parameters:,} ({format_params(ckpt_info.n_parameters)})")
+            _console.print(f"[bold]Tensors:[/bold] {ckpt_info.n_tensors}")
+            _console.print()
+
+            table = Table(title="Tensors", show_lines=False)
+            table.add_column("Name", style="cyan")
+            table.add_column("Shape", justify="right")
+            table.add_column("DType", justify="center")
+            table.add_column("Params", justify="right")
+            table.add_column("Size", justify="right")
+
+            for t in ckpt_info.tensors[:50]:
+                table.add_row(
+                    t.name, t.shape_str, t.dtype.value,
+                    f"{t.numel:,}", format_size(t.size_bytes),
+                )
+            if len(ckpt_info.tensors) > 50:
+                table.add_row("...", f"({len(ckpt_info.tensors) - 50} more)", "", "", "")
+
+            _console.print(table)
+        else:
+            click.echo(f"File: {ckpt_info.path}")
+            click.echo(f"Format: {ckpt_info.format.value}")
+            click.echo(f"Parameters: {ckpt_info.n_parameters:,}")
+            for t in ckpt_info.tensors[:50]:
+                click.echo(f"  {t.name}: {t.shape} {t.dtype.value}")
+
+    @cli.command(name="diff")
+    @click.argument("path_a", type=click.Path(exists=True))
+    @click.argument("path_b", type=click.Path(exists=True))
+    def diff_cmd(path_a: str, path_b: str) -> None:
+        """Compare two checkpoint files."""
+        result = diff(path_a, path_b)
+
+        if _HAS_RICH and _console is not None:
+            _console.print(f"\n[bold]A:[/bold] {result.path_a}")
+            _console.print(f"[bold]B:[/bold] {result.path_b}")
+            _console.print(f"Shared: {result.n_shared}  Identical: {result.n_identical}  Changes: {result.n_changes}")
+            _console.print()
+
+            if result.entries:
+                table = Table(show_lines=False)
+                table.add_column("", style="bold", width=3)
+                table.add_column("Tensor", style="cyan")
+                table.add_column("Change")
+                table.add_column("Details")
+
+                for e in result.entries:
+                    symbols = {"added": "[green]+[/green]", "removed": "[red]-[/red]",
+                             "shape_changed": "[yellow]~[/yellow]", "dtype_changed": "[yellow]~[/yellow]"}
+                    table.add_row(symbols.get(e.change_type, "?"), e.tensor_name, e.change_type, e.details)
+                _console.print(table)
+            else:
+                _console.print("[green]Checkpoints are structurally identical.[/green]")
+        else:
+            click.echo(format_diff(result))
+
+    @cli.command()
+    @click.argument("path", type=click.Path(exists=True))
+    def stats(path: str) -> None:
+        """Show checkpoint statistics."""
+        ckpt_info = inspect(path)
+        ckpt_stats = stats_from_info(ckpt_info)
+        click.echo(format_stats(ckpt_stats))
+
+    @cli.command(name="validate")
+    @click.argument("path", type=click.Path(exists=True))
+    def validate_cmd(path: str) -> None:
+        """Validate checkpoint integrity."""
+        result = validate(path)
+        if result.valid:
+            click.echo(f"✓ {path}: valid ({result.format.value})")
+        else:
+            click.echo(f"✗ {path}: invalid")
+        for issue in result.issues:
+            prefix = "  ERROR:" if issue.severity == "error" else "  WARN:"
+            click.echo(f"{prefix} {issue.message}")
+        if not result.valid:
+            raise SystemExit(1)
+
+    return cli
+
+
+cli = _build_cli()
+
+
+def main() -> None:
+    if cli is None:
+        print(
+            "The CLI requires extra dependencies. Install with:\n"
+            "  pip install ckptkit[cli]",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    cli()
+
+
+if __name__ == "__main__":
+    main()