main.py

#!/usr/bin/env python3

import argparse
import json
import os
import sys

from model_reader import read_model
from imatrix_reader import read_imatrix, detect_tied_groups, build_importance_table
from classifier import optimal_classify, compute_stats
from config_generator import generate_flags, format_flags
from quantizer import run_dry_run, run_quantization
from constants import CLASS_HARD_FLOORS


def _get_base_type(model: dict) -> str:
    is_qat = model.get("features", {}).get("is_qat", False)
    return "IQ4_XS" if is_qat else "Q5_K_M"


def main():
    parser = argparse.ArgumentParser(
        description="SHQ-program: imatrix-driven hybrid quantization"
    )
    parser.add_argument("--model", help="BF16 GGUF model path")
    parser.add_argument("--imatrix", action="append", default=[],
                        help="Imatrix GGUF path (can be specified multiple times)")
    parser.add_argument("--imatrix-method", choices=["max", "mean"], default="max",
                        help="How to combine multiple imatrix: max (conservative) or mean (default: max)")
    parser.add_argument("--size", type=float, default=6800,
                        help="Target file size in MiB (default: 6800 = ~6.6 GB)")
    parser.add_argument("--output", default=None, help="Output GGUF path")
    parser.add_argument("--run", action="store_true", help="Execute quantization")
    parser.add_argument("--show-config", action="store_true", help="Print config and exit")
    parser.add_argument("--verbose", action="store_true", help="Detailed output")
    parser.add_argument("--allow-q3-or-lower", action="store_true",
                        help="Allow Q3_K for low-importance tensors (risk of quality loss)")
    parser.add_argument("--aggro", type=float, default=None,
                        help="[deprecated] Use --size instead")
    parser.add_argument("--show-floors", action="store_true",
                        help="Print class hard floors and exit")

    args = parser.parse_args()

    if args.show_floors:
        _show_floors()
        return

    if not args.model or not args.imatrix:
        parser.print_usage()
        print("main.py: error: --model and --imatrix are required")
        sys.exit(1)

    target_mib = args.size

    print("=== SHQ-program ===")
    print(f"Model:   {args.model}")
    if len(args.imatrix) == 1:
        print(f"Imatrix: {args.imatrix[0]}")
    else:
        print(f"Imatrix: {len(args.imatrix)} files ({args.imatrix_method})")
        for p in args.imatrix:
            print(f"  - {p}")
    print(f"Target:  {target_mib:.0f} MiB ({target_mib / 1024:.2f} GB)")
    if args.allow_q3_or_lower:
        print("  --allow-q3-or-lower: low-importance tensors may go to Q3_K")
    print()

    print("[1/4] Reading model...")
    model = read_model(args.model)
    print(f"  Architecture: {model['architecture']}")
    print(f"  Tensors: {model['n_tensors']}")
    print(f"  Features: {json.dumps(model['features'], indent=2)}")

    print("\n[2/4] Reading imatrix...")
    imatrix_list = [read_imatrix(p) for p in args.imatrix]
    for im in imatrix_list:
        print(f"  {im['path']}: {im['n_tensors']} tensors, datasets={im['meta'].get('imatrix.datasets', '?')}")

    from imatrix_reader import combine_imatrix
    imatrix = combine_imatrix(imatrix_list, method=args.imatrix_method)
    print(f"  Combined: {imatrix['n_tensors']} tensors")

    print("\n[3/4] Detecting tied groups...")
    tied_groups = detect_tied_groups(imatrix)
    print(f"  Found {len(tied_groups)} tied groups:")
    for g in tied_groups:
        if len(g) > 1:
            print(f"    TIED ({len(g)}): {g[0].replace('.weight', '')}  =  "
                  f"{g[1].replace('.weight', '')}")

    imp_table = build_importance_table(imatrix, model)

    print("\n[4/4] Classifying tensors (greedy imatrix-driven)...")

    # Получаем и маппинг тиров, и точную карту паддингов напрямую из классификатора
    assignments, padded_ne_map = optimal_classify(
        imp_table, tied_groups, model,
        target_size_mib=target_mib,
        allow_q3=args.allow_q3_or_lower,
    )
    
    ne_map = {k: v["n_elements"] for k, v in model.get("tensors", {}).items()}
    for tname, info in imp_table.items():
        if tname not in ne_map:
            ne_map[tname] = info["n_elements"]

    # Передаем padded_ne_map для корректного вывода логов на экран
    _show_tier_summary(assignments, imp_table, ne_map, padded_ne_map)

    base_type = _get_base_type(model)
    flags = generate_flags(assignments, model, base_type, target_mib)
    flags["imatrix"] = args.imatrix

    print(f"\nConfig (base={flags['base_type']}):")
    print(format_flags(flags))

    if args.show_config:
        return

    print("\n--- Dry Run ---")
    dry_size = run_dry_run(flags, args.model)
    _show_size_result(dry_size, target_mib)

    if not args.run:
        print("\nDry run only. Use --run to execute quantization.")
        return

    if not args.output:
        base = os.path.splitext(os.path.basename(args.model))[0]
        args.output = base + "-SHQ.gguf"

    print(f"\n--- Running quantization: {args.output} ---")
    success = run_quantization(flags, args.model, args.output)
    if success:
        print("Done!")
    else:
        print("Failed!")
        sys.exit(1)


def _show_tier_summary(assignments, imp_table, ne_map, padded_ne_map=None):
    stats = compute_stats(assignments, ne_map, padded_ne_map)
    
    print("\n  Tier distribution:")
    for tier in sorted(stats["by_tier_count"].keys()):
        count = stats["by_tier_count"][tier]
        mib = stats["by_tier_mib"].get(tier, 0.0)
        print(f"    {tier}: {count} tensors ({mib:.1f} MiB)")
    print(f"  Total estimated size: {stats['total_mib']:.1f} MiB")

    ranked = sorted(
        [(n, v) for n, v in imp_table.items()],
        key=lambda x: -x[1]["importance_mean"],
    )
    print("\n  Top 10 by importance:")
    for n, v in ranked[:10]:
        tier = assignments.get(n, "base")
        display = n.replace(".weight", "").replace(".bias", "")
        print(f"    {display[:52]:52s} imp={v['importance_mean']:10.0f} tier={tier}")


def _show_size_result(dry_size, target_mib):
    if dry_size:
        print(f"  Estimated size: {dry_size:.0f} MiB ({dry_size / 1024:.2f} GB)")
        diff = dry_size - target_mib
        if diff > 0:
            print(f"  ⚠  Over target by {diff:.0f} MiB")
        else:
            print(f"  ✓ Under target by {-diff:.0f} MiB")
    else:
        print("  ⚠  Could not parse size from dry-run output")


def _show_floors():
    print("  Class hard floors (never below without --allow-q3-or-lower):\n")
    max_n = max(len(c) for c in CLASS_HARD_FLOORS)
    for cls, floor in sorted(CLASS_HARD_FLOORS.items()):
        print(f"    {cls:<{max_n}}  →  {floor}")
    print(f"\n  Default floor (unknown class): Q4_K")
    print(f"  --allow-q3-or-lower enables Q3_K for: ffn_down, attn_output, ssm_out")


if __name__ == "__main__":
    main()