#!/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()