File size: 3,947 Bytes
628e2fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """Compress a verified GLiNER2.5-Decide classification Core ML model."""
import argparse
import json
import time
from pathlib import Path
import coremltools as ct
from coremltools.optimize.coreml import (
OpLinearQuantizerConfig,
OpPalettizerConfig,
OptimizationConfig,
get_weights_metadata,
linear_quantize_weights,
palettize_weights,
)
def package_bytes(path):
return sum(item.stat().st_size for item in path.rglob("*") if item.is_file())
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--source", required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--bits", type=int, choices=[4, 6, 8], default=8)
parser.add_argument("--method", choices=["lut", "linear", "embedding"], default="lut")
parser.add_argument("--linear-mode", choices=["linear_symmetric", "linear"], default="linear_symmetric")
parser.add_argument("--granularity", choices=["per_tensor", "per_grouped_channel"], default="per_tensor")
parser.add_argument("--group-size", type=int, default=32)
args = parser.parse_args()
start = time.perf_counter()
model = ct.models.MLModel(args.source, skip_model_load=True)
selected_weights = None
if args.method in ("linear", "embedding"):
if args.bits != 8:
parser.error("linear quantization here supports only 8-bit weights")
quantizer = OpLinearQuantizerConfig(
mode=args.linear_mode, dtype="int8", granularity="per_channel", weight_threshold=2048,
)
if args.method == "embedding":
metadata = get_weights_metadata(model, weight_threshold=2048)
names = [
name for name, weight in metadata.items()
if "encoder_embeddings_word_embeddings" in name
and len(weight.val.shape) == 2
and any(op.op_type == "gather" for op in weight.child_ops)
]
if len(names) != 1:
raise ValueError(f"expected one token embedding weight, found {names}")
config = OptimizationConfig(op_name_configs={names[0]: quantizer})
selected_weights = names
else:
config = OptimizationConfig(global_config=quantizer)
compressed = linear_quantize_weights(model, config=config)
variant = (
f"W8 {args.linear_mode} token embedding only"
if args.method == "embedding" else f"W8 {args.linear_mode} per-channel"
)
else:
config = OptimizationConfig(global_config=OpPalettizerConfig(
mode="kmeans", nbits=args.bits, granularity=args.granularity,
group_size=args.group_size,
enable_per_channel_scale=args.granularity == "per_grouped_channel",
num_kmeans_workers=4,
))
compressed = palettize_weights(model, config=config)
variant = f"LUT{args.bits} {args.granularity}"
compressed.short_description = f"{model.short_description}; {variant} weights"
compressed.author = model.author
compressed.license = model.license
compressed.user_defined_metadata.update(model.user_defined_metadata)
compressed.user_defined_metadata["weight_compression"] = variant
compressed.save(args.output)
report = {"source": args.source, "output": args.output, "method": args.method,
"linear_mode": args.linear_mode if args.method in ("linear", "embedding") else None,
"granularity": "per_channel" if args.method in ("linear", "embedding") else args.granularity,
"selected_weights": selected_weights,
"source_bytes": package_bytes(Path(args.source)),
"output_bytes": package_bytes(Path(args.output)), "compression_seconds": time.perf_counter() - start}
Path(args.output).with_suffix(".json").write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|