OBLITERATUS/obliteratus/restore_multimodal.py

#!/usr/bin/env python3
"""Restore multimodal components (vision tower, MTP head, etc.) after text-only abliteration.

OBLITERATUS surgery strips models to text-only weights with simplified tensor
names (model.layers.* instead of model.language_model.layers.*). This works
for text inference but breaks GGUF conversion (missing MTP block) and drops
the vision tower entirely.

This tool merges abliterated text weights back into the full stock model,
restoring all non-text components from stock while keeping the abliterated
text weights intact. The merge is lossless — no blending or interpolation.

Usage:
    obliteratus restore-multimodal \\
        --abliterated outputs/my-abliterated-model \\
        --stock Qwen/Qwen3.8-27B \\
        --output outputs/my-abliterated-model-full

Typical tensor counts for Qwen3.8-27B:
    Abliterated (text-only):  851 tensors (model.layers.*)
    Stock (full multimodal): 1199 tensors (model.language_model.*, mtp.*, visual.*)
    Merged output:           1199 tensors (abliterated text + stock vision/MTP)
"""

from __future__ import annotations

import json
import logging
import shutil
from pathlib import Path

from safetensors.torch import load_file, save_file

logger = logging.getLogger(__name__)

# Common prefix mappings: text-only key -> full multimodal key
_PREFIX_STRIPS = [
    ("model.language_model.", "model."),
]


def _normalize_key(key: str) -> str:
    """Strip multimodal prefixes to match text-only tensor names."""
    for full_prefix, stripped_prefix in _PREFIX_STRIPS:
        if key.startswith(full_prefix):
            return stripped_prefix + key[len(full_prefix):]
    return key


def _build_abliterated_lookup(abl_dir: Path) -> dict:
    """Build a lookup from tensor key -> (shard_file, original_key)."""
    index_path = abl_dir / "model.safetensors.index.json"
    if not index_path.exists():
        # Single-file model
        sf_files = list(abl_dir.glob("*.safetensors"))
        if len(sf_files) == 1:
            tensors = load_file(str(sf_files[0]))
            return {k: (sf_files[0].name, k) for k in tensors}
        raise FileNotFoundError(f"No safetensors index or single file in {abl_dir}")

    with open(index_path) as f:
        index = json.load(f)
    return {key: (shard, key) for key, shard in index["weight_map"].items()}


def restore_multimodal(
    abliterated_dir: str | Path,
    stock_dir: str | Path,
    output_dir: str | Path,
    skip_existing: bool = True,
) -> dict:
    """Merge abliterated text weights into full stock model.

    Args:
        abliterated_dir: Path to text-only abliterated model.
        stock_dir: Path to stock (full multimodal) model.
        output_dir: Path to write merged output.
        skip_existing: Skip files that already exist in output_dir.

    Returns:
        dict with counts: replaced, kept, total.
    """
    abl_dir = Path(abliterated_dir)
    stock_dir = Path(stock_dir)
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # 1. Copy all non-safetensor files from stock
    skip_names = {".git", "__pycache__", "crc32.txt", "README.md", ".gitattributes"}
    for f in stock_dir.iterdir():
        if f.name.endswith(".safetensors") or f.name == "model.safetensors.index.json":
            continue
        if f.name in skip_names:
            continue
        dst = out_dir / f.name
        if not dst.exists() or not skip_existing:
            shutil.copy2(f, dst)
            logger.info("Copied %s", f.name)

    # Copy abliteration metadata if present
    for extra in ["abliteration_metadata.json", "hard_negative_residue.json"]:
        src = abl_dir / extra
        if src.exists():
            shutil.copy2(src, out_dir / extra)
            logger.info("Copied %s", extra)

    # 2. Build abliterated tensor lookup
    abl_keys = _build_abliterated_lookup(abl_dir)
    logger.info("Abliterated model: %d tensors", len(abl_keys))

    # 3. Load stock index
    with open(stock_dir / "model.safetensors.index.json") as f:
        stock_index = json.load(f)

    stock_shards = sorted(set(stock_index["weight_map"].values()))
    logger.info("Stock model: %d tensors across %d shards",
                len(stock_index["weight_map"]), len(stock_shards))

    # 4. Process each stock shard — replace text tensors with abliterated versions
    abl_shard_cache = {}
    replaced = 0
    kept = 0
    new_weight_map = {}

    for shard_idx, stock_shard in enumerate(stock_shards):
        logger.info("[%d/%d] %s", shard_idx + 1, len(stock_shards), stock_shard)

        stock_tensors = load_file(str(stock_dir / stock_shard))
        merged = {}

        for key, tensor in stock_tensors.items():
            norm_key = _normalize_key(key)

            # Try normalized key first, then direct match
            match_key = norm_key if norm_key in abl_keys else (key if key in abl_keys else None)

            if match_key is not None:
                abl_shard, abl_key = abl_keys[match_key]
                if abl_shard not in abl_shard_cache:
                    abl_shard_cache[abl_shard] = load_file(str(abl_dir / abl_shard))
                merged[key] = abl_shard_cache[abl_shard][abl_key]
                replaced += 1
            else:
                merged[key] = tensor
                kept += 1

            new_weight_map[key] = stock_shard

        save_file(merged, str(out_dir / stock_shard))
        logger.info("  Saved %s (%d tensors)", stock_shard, len(merged))

        # Free cache periodically
        if len(abl_shard_cache) > 5:
            abl_shard_cache.clear()

    # 5. Write merged index
    new_index = {
        "metadata": stock_index.get("metadata", {}),
        "weight_map": new_weight_map,
    }
    with open(out_dir / "model.safetensors.index.json", "w") as f:
        json.dump(new_index, f, indent=2)

    result = {"replaced": replaced, "kept": kept, "total": replaced + kept}
    logger.info("Replaced: %d tensors (abliterated)", replaced)
    logger.info("Kept:     %d tensors (from stock)", kept)
    logger.info("Output:   %s", out_dir)
    return result


def main():
    import argparse
    logging.basicConfig(level=logging.INFO, format="%(message)s")

    p = argparse.ArgumentParser(
        description="Restore vision tower, MTP head, and other multimodal "
                    "components after text-only abliteration surgery."
    )
    p.add_argument("--abliterated", required=True,
                   help="Path to text-only abliterated model directory")
    p.add_argument("--stock", required=True,
                   help="Path to stock (full) model directory or HF repo ID")
    p.add_argument("--output", required=True,
                   help="Output directory for merged model")
    args = p.parse_args()

    result = restore_multimodal(args.abliterated, args.stock, args.output)
    print(f"\nDone: {result['replaced']} abliterated + {result['kept']} stock = {result['total']} total")


if __name__ == "__main__":
    main()