| """Public, reusable pruning interface for Hugging Face model repositories.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import asdict, replace |
| from pathlib import Path |
| from typing import Any, Callable, Iterable, Sequence |
|
|
| from huggingface_hub import hf_hub_download |
| from torch import nn |
|
|
| from .core import PruneConfig, PruneReport, prune |
| from .hf import prune_hf_model |
|
|
|
|
| class HigherOrderPruner: |
| """Configuration-backed API for higher-order pruning. |
| |
| The object contains method defaults, not model weights. ``from_pretrained`` |
| reads those defaults from a Hugging Face *model* repository. The same |
| instance can prune any compatible model; it is not tied to Qwen or Falcon. |
| """ |
|
|
| CONFIG_NAME = "pruner_config.json" |
|
|
| def __init__(self, config: PruneConfig | None = None): |
| self.config = config or PruneConfig(sparsity=0.5) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| repo_id: str, |
| *, |
| revision: str = "main", |
| token: str | bool | None = None, |
| cache_dir: str | Path | None = None, |
| local_files_only: bool = False, |
| ) -> "HigherOrderPruner": |
| path = hf_hub_download( |
| repo_id=repo_id, |
| filename=cls.CONFIG_NAME, |
| revision=revision, |
| token=token, |
| cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| ) |
| payload = json.loads(Path(path).read_text()) |
| if payload.get("format") != "higher-order-pruner-v1": |
| raise ValueError("Unsupported pruner configuration format") |
| values = dict(payload["defaults"]) |
| if "backtracking" in values: |
| values["backtracking"] = tuple(values["backtracking"]) |
| return cls(PruneConfig(**values)) |
|
|
| def save_pretrained(self, directory: str | Path) -> Path: |
| destination = Path(directory) |
| destination.mkdir(parents=True, exist_ok=True) |
| path = destination / self.CONFIG_NAME |
| path.write_text(json.dumps({ |
| "format": "higher-order-pruner-v1", |
| "method": "diagonal-higher-order-kfac", |
| "defaults": asdict(self.config), |
| }, indent=2) + "\n") |
| return path |
|
|
| def prune( |
| self, |
| model: nn.Module, |
| calibration_batches: Iterable[Any], |
| loss_fn: Callable[[nn.Module, Any], Any], |
| *, |
| sparsity: float | None = None, |
| validation_batches: Iterable[Any] | None = None, |
| module_names: Sequence[str] | None = None, |
| progress_dir: str | Path | None = None, |
| on_progress: Callable[[dict[str, Any]], None] | None = None, |
| ) -> PruneReport: |
| config = self.config if sparsity is None else replace(self.config, sparsity=sparsity) |
| return prune( |
| model, |
| calibration_batches, |
| loss_fn, |
| config, |
| validation_batches=validation_batches, |
| module_names=module_names, |
| progress_dir=progress_dir, |
| on_progress=on_progress, |
| ) |
|
|
| def prune_hf_model( |
| self, |
| model_id: str, |
| calibration_texts: Sequence[str], |
| *, |
| output_dir: str | Path, |
| sparsity: float | None = None, |
| validation_texts: Sequence[str] | None = None, |
| revision: str = "main", |
| sequence_length: int = 256, |
| module_names: Sequence[str] | None = None, |
| device: str = "cuda", |
| on_progress: Callable[[dict], None] | None = None, |
| ) -> dict: |
| selected = self.config.sparsity if sparsity is None else sparsity |
| return prune_hf_model( |
| model_id, |
| calibration_texts, |
| validation_texts=validation_texts, |
| sparsity=selected, |
| output_dir=output_dir, |
| revision=revision, |
| sequence_length=sequence_length, |
| probes=self.config.probes, |
| module_names=module_names, |
| device=device, |
| on_progress=on_progress, |
| ) |
|
|
