"""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 @dataclass 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