Spaces:
Paused
Paused
| """Publish a rank-1 patch to the Hugging Face Hub. | |
| Only the shard that holds the edited tensors is downloaded. Untouched files are | |
| copied on the Hub with CommitOperationCopy when the source and target repos match. | |
| A write token is used for that commit and is not written to disk by this module. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import io | |
| import json | |
| import tempfile | |
| from dataclasses import asdict, dataclass | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Sequence, Tuple | |
| import numpy as np | |
| from modelifier.engine import audit, rank1_column_ablation, resolve_tensor | |
| class HubPatchAudit: | |
| layer_idx: int | |
| tensor_name: str | |
| shard_file: str | |
| touched_columns: int | |
| touched_elements: int | |
| relative_frobenius_pct: float | |
| sigma2_to_sigma1_pct: float | |
| component_scale: float | |
| sha256_before: str | |
| sha256_after: str | |
| def _sha256(blob: bytes) -> str: | |
| return hashlib.sha256(blob).hexdigest() | |
| def _to_numpy(tensor) -> np.ndarray: | |
| if isinstance(tensor, np.ndarray): | |
| array = tensor | |
| else: | |
| array = tensor.detach().to(dtype=__import__("torch").float32).cpu().numpy() | |
| if array.ndim != 2: | |
| raise ValueError(f"Expected a matrix, got {array.shape}.") | |
| return np.ascontiguousarray(array, dtype=np.float32) | |
| def _as_target(module_suffix: str): | |
| from modelifier.engine import SemanticTarget | |
| if module_suffix in {item.value for item in SemanticTarget}: | |
| return SemanticTarget(module_suffix) | |
| return module_suffix | |
| def inspect_repo( | |
| source_repo_id: str, | |
| layers: Sequence[int], | |
| module_suffix: str = "residual_out_proj", | |
| token: Optional[str] = None, | |
| ) -> str: | |
| """Read only the weight index (or one small header) and say which shard would change.""" | |
| from huggingface_hub import hf_hub_download, list_repo_files | |
| files = list(list_repo_files(source_repo_id, token=token)) | |
| lines = [f"Repository `{source_repo_id}` has {len(files)} files."] | |
| with tempfile.TemporaryDirectory() as tmp: | |
| tmp_dir = Path(tmp) | |
| if "model.safetensors.index.json" in files: | |
| local = hf_hub_download( | |
| repo_id=source_repo_id, | |
| filename="model.safetensors.index.json", | |
| local_dir=tmp_dir, | |
| token=token, | |
| ) | |
| weight_map = json.loads(Path(local).read_text(encoding="utf-8"))["weight_map"] | |
| elif "model.safetensors" in files: | |
| weight_map = None | |
| lines.append("Single file: `model.safetensors` holds every tensor.") | |
| else: | |
| raise FileNotFoundError("No model.safetensors weights in that repository.") | |
| if weight_map is not None: | |
| keys = list(weight_map) | |
| for layer in layers: | |
| try: | |
| key = resolve_tensor(keys, int(layer), _as_target(module_suffix)) | |
| except (KeyError, ValueError) as exc: | |
| lines.append(f"- layer {layer}: not resolved ({exc})") | |
| continue | |
| lines.append(f"- layer {layer}: `{key}` in `{weight_map[key]}`") | |
| else: | |
| lines.append("Pass a sharded model to see the per-layer shard map. This one is a single file.") | |
| lines.append("Nothing was modified.") | |
| return "\n".join(lines) | |
| def patch_and_publish( | |
| source_repo_id: str, | |
| target_repo_id: str, | |
| layer_directions: Dict[int, np.ndarray], | |
| top_k: int = 64, | |
| alpha: float = 6.0, | |
| module_suffix: str = "residual_out_proj", | |
| token: Optional[str] = None, | |
| private: bool = False, | |
| ) -> Tuple[str, List[HubPatchAudit]]: | |
| import torch | |
| from huggingface_hub import ( | |
| CommitOperationAdd, | |
| CommitOperationCopy, | |
| HfApi, | |
| hf_hub_download, | |
| ) | |
| from safetensors import safe_open | |
| from safetensors.torch import save_file | |
| api = HfApi(token=token) | |
| info = api.model_info(source_repo_id) | |
| all_files = [s.rfilename for s in info.siblings] | |
| audits: List[HubPatchAudit] = [] | |
| with tempfile.TemporaryDirectory() as tmp: | |
| tmp_dir = Path(tmp) | |
| if "model.safetensors.index.json" in all_files: | |
| idx = hf_hub_download( | |
| repo_id=source_repo_id, | |
| filename="model.safetensors.index.json", | |
| local_dir=tmp_dir, | |
| token=token, | |
| ) | |
| weight_map: Dict[str, str] = json.loads(Path(idx).read_text(encoding="utf-8"))["weight_map"] | |
| keys = list(weight_map) | |
| elif "model.safetensors" in all_files: | |
| local = hf_hub_download( | |
| repo_id=source_repo_id, | |
| filename="model.safetensors", | |
| local_dir=tmp_dir, | |
| token=token, | |
| ) | |
| with safe_open(local, framework="pt", device="cpu") as handle: | |
| keys = list(handle.keys()) | |
| weight_map = {k: "model.safetensors" for k in keys} | |
| else: | |
| raise FileNotFoundError(f"No safetensors in {source_repo_id}.") | |
| grouped: Dict[str, List[Tuple[int, str, np.ndarray]]] = {} | |
| for layer, direction in layer_directions.items(): | |
| key = resolve_tensor(keys, int(layer), _as_target(module_suffix)) | |
| grouped.setdefault(weight_map[key], []).append((int(layer), key, np.asarray(direction, dtype=np.float32))) | |
| patched: Dict[str, Path] = {} | |
| for shard, edits in grouped.items(): | |
| local_shard = Path( | |
| hf_hub_download(repo_id=source_repo_id, filename=shard, local_dir=tmp_dir, token=token) | |
| ) | |
| tensors = {} | |
| with safe_open(str(local_shard), framework="pt", device="cpu") as handle: | |
| metadata = handle.metadata() | |
| for key in handle.keys(): | |
| tensors[key] = handle.get_tensor(key) | |
| for layer, key, direction in edits: | |
| original = tensors[key] | |
| before = _to_numpy(original) | |
| raw_before = original.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() | |
| updated, cols = rank1_column_ablation(before, direction, top_k=top_k, alpha=alpha) | |
| native = torch.from_numpy(updated).to(dtype=original.dtype) | |
| tensors[key] = native | |
| raw_after = native.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() | |
| stats = audit(before, _to_numpy(native)) | |
| u = direction.reshape(-1).astype(np.float32) | |
| u = u / max(float(np.linalg.norm(u)), 1e-15) | |
| work = before if before.shape[0] == u.shape[0] else before.T | |
| after_work = updated if updated.shape[0] == u.shape[0] else updated.T | |
| b = float((u @ work)[cols].mean()) | |
| a = float((u @ after_work)[cols].mean()) | |
| audits.append( | |
| HubPatchAudit( | |
| layer_idx=layer, | |
| tensor_name=key, | |
| shard_file=shard, | |
| touched_columns=int(len(cols)), | |
| touched_elements=int(work.shape[0] * len(cols)), | |
| relative_frobenius_pct=round(stats["relative_frobenius_pct"], 4), | |
| sigma2_to_sigma1_pct=round(stats["sigma2_over_sigma1_pct"], 4), | |
| component_scale=round(a / b, 4) if abs(b) > 1e-8 else 0.0, | |
| sha256_before=_sha256(raw_before), | |
| sha256_after=_sha256(raw_after), | |
| ) | |
| ) | |
| out = tmp_dir / f"patched_{Path(shard).name}" | |
| save_file(tensors, str(out), metadata=metadata) | |
| patched[shard] = out | |
| api.create_repo(repo_id=target_repo_id, private=private, exist_ok=True, repo_type="model") | |
| operations = [ | |
| CommitOperationAdd(path_in_repo=shard, path_or_fileobj=str(path)) | |
| for shard, path in patched.items() | |
| ] | |
| for name in all_files: | |
| if name in patched or name.startswith(".git"): | |
| continue | |
| if source_repo_id == target_repo_id: | |
| operations.append(CommitOperationCopy(src_path_in_repo=name, path_in_repo=name)) | |
| elif not name.endswith((".safetensors", ".bin", ".pt", ".gguf", ".onnx")): | |
| local = hf_hub_download(repo_id=source_repo_id, filename=name, local_dir=tmp_dir, token=token) | |
| operations.append(CommitOperationAdd(path_in_repo=name, path_or_fileobj=local)) | |
| manifest = { | |
| "source_model": source_repo_id, | |
| "target_model": target_repo_id, | |
| "alpha": alpha, | |
| "top_k_columns": top_k, | |
| "formula": "w_j <- w_j - alpha * (u^T w_j) * u", | |
| "note": "alpha 1 removes the component. alpha 6, the published Qwen strength, multiplies it by -5.", | |
| "audits": [asdict(item) for item in audits], | |
| } | |
| operations.append( | |
| CommitOperationAdd( | |
| path_in_repo="weight_space_manifest.json", | |
| path_or_fileobj=io.BytesIO(json.dumps(manifest, indent=2).encode("utf-8")), | |
| ) | |
| ) | |
| commit = api.create_commit( | |
| repo_id=target_repo_id, | |
| repo_type="model", | |
| operations=operations, | |
| commit_message=f"Rank-1 weight-space patch, alpha={alpha}, layers={list(layer_directions)}", | |
| ) | |
| return commit.commit_url, audits | |