utils.py

import numpy as np
from constants import GGUF_TYPE_NAMES, GGUF_TYPE_NAMES_INV


def type_name(type_id: int) -> str:
    return GGUF_TYPE_NAMES.get(type_id, f"UNKNOWN({type_id})")


def type_id(name: str) -> int:
    return GGUF_TYPE_NAMES_INV.get(name, -1)


def parse_size_line(line: str) -> float | None:
    import re
    m = re.search(r"quant size\s*=\s*([0-9.]+)\s*MiB", line)
    if m:
        return float(m.group(1))
    m = re.search(r"model size\s*=\s*([0-9.]+)\s*MiB", line)
    if m:
        return float(m.group(1))
    return None


def parse_quant_size(output: str) -> float | None:
    """Extract quant size from full output (prefers quant over model size)."""
    import re
    # Try quant size first across all lines
    m = re.search(r"quant size\s*=\s*([0-9.]+)\s*MiB", output)
    if m:
        return float(m.group(1))
    # Fall back to model size
    m = re.search(r"model size\s*=\s*([0-9.]+)\s*MiB", output)
    if m:
        return float(m.group(1))
    return None


def parse_fallback_warnings(output: str) -> int:
    import re
    return len(re.findall(r"converting to\s+(q[0-9]_[0-9KMS]|iq[0-9])", output))


def format_size(mib: float) -> str:
    if mib >= 1024:
        return f"{mib/1024:.2f} GB"
    return f"{mib:.0f} MiB"


def get_tensor_type(tensor_name: str) -> str:
    parts = tensor_name.split(".")
    if len(parts) >= 2 and parts[0] == "blk":
        return parts[2] if len(parts) >= 3 else "unknown"
    return "global"