diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 9de29aa..bed2423 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -2,15 +2,24 @@ "schema_version": 2, "owner": "OBLITERATUS maintainers", "source_inventory": { - "roots": ["app.py", "obliteratus"], - "coverage_roots": ["obliteratus"] + "roots": [ + "app.py", + "obliteratus" + ], + "coverage_roots": [ + "obliteratus" + ] }, "contract_surfaces": [ { "id": "package-entrypoints", "owner": "packaging maintainers", "description": "Installed package exports, module entry points, console dispatch, and optional application construction", - "contract_types": ["package-entrypoint", "public-interface", "operator-ui"], + "contract_types": [ + "package-entrypoint", + "public-interface", + "operator-ui" + ], "paths": [ "app.py", "obliteratus/__init__.py", @@ -33,7 +42,12 @@ "id": "core-pipeline", "owner": "abliteration pipeline maintainers", "description": "Model mutation, orchestration, automated search, study execution, sweeps, and tournament lifecycle", - "contract_types": ["model-mutation", "model-runtime", "orchestration", "persistence"], + "contract_types": [ + "model-mutation", + "model-runtime", + "orchestration", + "persistence" + ], "paths": [ "obliteratus/abliterate.py", "obliteratus/persistence_contracts.py", @@ -44,7 +58,8 @@ "obliteratus/runner.py", "obliteratus/sweep.py", "obliteratus/tourney.py", - "obliteratus/tourney_contracts.py" + "obliteratus/tourney_contracts.py", + "obliteratus/restore_multimodal.py" ], "required_tests": [ "tests/test_abliterate.py", @@ -66,7 +81,12 @@ "id": "configuration-and-selection", "owner": "runtime compatibility maintainers", "description": "Configuration, defaults, model profiles, architecture detection, device selection, loading, and preset contracts", - "contract_types": ["configuration", "architecture-selection", "device-boundary", "model-runtime"], + "contract_types": [ + "configuration", + "architecture-selection", + "device-boundary", + "model-runtime" + ], "paths": [ "obliteratus/adaptive_defaults.py", "obliteratus/architecture_profiles.py", @@ -99,7 +119,12 @@ "id": "public-interface-and-services", "owner": "operator interface maintainers", "description": "CLI and interactive interfaces, remote execution, service adapters, scheduling, and UI boundaries", - "contract_types": ["public-interface", "external-service", "remote-execution", "operator-ui"], + "contract_types": [ + "public-interface", + "external-service", + "remote-execution", + "operator-ui" + ], "paths": [ "obliteratus/bestiary_sync.py", "obliteratus/cli.py", @@ -133,7 +158,12 @@ "id": "research-inputs-and-provenance", "owner": "research integrity maintainers", "description": "Prompt sources, hard negatives, community records, reproducibility controls, and telemetry schemas", - "contract_types": ["research-input", "research-output", "reproducibility", "persistence"], + "contract_types": [ + "research-input", + "research-output", + "reproducibility", + "persistence" + ], "paths": [ "obliteratus/community.py", "obliteratus/hard_negative.py", @@ -154,7 +184,12 @@ "id": "analysis-methods", "owner": "mechanistic analysis maintainers", "description": "Mechanistic analysis algorithms, numerical invariants, transfer methods, and visualization contracts", - "contract_types": ["research-metric", "numerical-invariant", "research-output", "model-runtime"], + "contract_types": [ + "research-metric", + "numerical-invariant", + "research-output", + "model-runtime" + ], "paths": [ "obliteratus/analysis/activation_patching.py", "obliteratus/analysis/activation_probing.py", @@ -205,7 +240,12 @@ "id": "evaluation-and-benchmarks", "owner": "evaluation maintainers", "description": "Evaluation routing, benchmark execution, plots, statistical metrics, and external adapter results", - "contract_types": ["research-metric", "numerical-invariant", "research-output", "external-service"], + "contract_types": [ + "research-metric", + "numerical-invariant", + "research-output", + "external-service" + ], "paths": [ "obliteratus/evaluation/advanced_metrics.py", "obliteratus/evaluation/baselines.py", @@ -233,8 +273,14 @@ "id": "reporting", "owner": "research reporting maintainers", "description": "Report schemas, deterministic serialization, plotting, and persisted result artifacts", - "contract_types": ["research-output", "persistence", "numerical-invariant"], - "paths": ["obliteratus/reporting/report.py"], + "contract_types": [ + "research-output", + "persistence", + "numerical-invariant" + ], + "paths": [ + "obliteratus/reporting/report.py" + ], "required_tests": [ "tests/test_report.py", "tests/test_evaluation_reporting_contracts.py", @@ -245,7 +291,11 @@ "id": "ablation-strategies", "owner": "ablation strategy maintainers", "description": "Strategy registration, enumeration, architecture navigation, mutation, and restoration behavior", - "contract_types": ["ablation-strategy", "architecture-selection", "model-mutation"], + "contract_types": [ + "ablation-strategy", + "architecture-selection", + "model-mutation" + ], "paths": [ "obliteratus/strategies/base.py", "obliteratus/strategies/embedding_ablation.py", @@ -268,28 +318,39 @@ "path": "obliteratus/cli.py", "risk_class": "cpu-contract", "risk": "public parsing, validation, dispatch, and local/remote option propagation", - "required_tests": ["tests/test_cli.py", "tests/test_cli_boundaries.py"], + "required_tests": [ + "tests/test_cli.py", + "tests/test_cli_boundaries.py" + ], "conditional_gates": [] }, { "path": "obliteratus/config.py", "risk_class": "cpu-contract", "risk": "configuration defaults, normalization, validation, and serialization", - "required_tests": ["tests/test_config.py", "tests/test_config_properties.py"], + "required_tests": [ + "tests/test_config.py", + "tests/test_config_properties.py" + ], "conditional_gates": [] }, { "path": "obliteratus/architecture_profiles.py", "risk_class": "cpu-contract", "risk": "architecture detection and projection-path contracts", - "required_tests": ["tests/test_architecture_profiles.py", "tests/test_gemma4_support.py"], + "required_tests": [ + "tests/test_architecture_profiles.py", + "tests/test_gemma4_support.py" + ], "conditional_gates": [] }, { "path": "obliteratus/community.py", "risk_class": "cpu-contract", "risk": "contribution schema, aggregation, and atomic persistence", - "required_tests": ["tests/test_community.py"], + "required_tests": [ + "tests/test_community.py" + ], "conditional_gates": [] }, { @@ -317,7 +378,10 @@ "path": "obliteratus/evaluation/metrics.py", "risk_class": "cpu-contract", "risk": "research metric semantics and numerical invariants", - "required_tests": ["tests/test_metrics.py", "tests/test_property_contracts.py"], + "required_tests": [ + "tests/test_metrics.py", + "tests/test_property_contracts.py" + ], "conditional_gates": [] }, { @@ -335,36 +399,57 @@ "path": "obliteratus/device.py", "risk_class": "mixed-runtime", "risk": "device and dtype selection across CPU, CUDA, and Apple backends", - "required_tests": ["tests/test_device_boundaries.py"], - "conditional_gates": ["cuda-runtime", "mps-runtime"] + "required_tests": [ + "tests/test_device_boundaries.py" + ], + "conditional_gates": [ + "cuda-runtime", + "mps-runtime" + ] }, { "path": "obliteratus/runtime_contracts.py", "risk_class": "cpu-contract", "risk": "model request, memory, architecture scale, and dtype capability decisions", - "required_tests": ["tests/test_runtime_contracts.py"], + "required_tests": [ + "tests/test_runtime_contracts.py" + ], "conditional_gates": [] }, { "path": "obliteratus/models/loader.py", "risk_class": "mixed-runtime", "risk": "model loading, cache, architecture, quantization, and device-map boundaries", - "required_tests": ["tests/test_loader_boundaries.py", "tests/test_offline_integration.py"], - "conditional_gates": ["cuda-runtime", "bitsandbytes-runtime"] + "required_tests": [ + "tests/test_loader_boundaries.py", + "tests/test_offline_integration.py" + ], + "conditional_gates": [ + "cuda-runtime", + "bitsandbytes-runtime" + ] }, { "path": "obliteratus/models/quant_dequant.py", "risk_class": "mixed-runtime", "risk": "FP8 and NVFP4 format detection, numerical dequantization, Hub revision integrity, and checkpoint materialization", - "required_tests": ["tests/test_quant_dequant.py"], - "conditional_gates": ["cuda-runtime"] + "required_tests": [ + "tests/test_quant_dequant.py" + ], + "conditional_gates": [ + "cuda-runtime" + ] }, { "path": "obliteratus/mlx_backend.py", "risk_class": "mixed-runtime", "risk": "MLX model discovery, tensor placement, mutation, and persistence", - "required_tests": ["tests/test_mlx_backend_boundaries.py"], - "conditional_gates": ["mlx-runtime"] + "required_tests": [ + "tests/test_mlx_backend_boundaries.py" + ], + "conditional_gates": [ + "mlx-runtime" + ] }, { "path": "obliteratus/persistence_contracts.py", @@ -390,14 +475,21 @@ "tests/test_projection_math_contracts.py", "tests/test_offline_integration.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/auto_obliterate.py", "risk_class": "mixed-runtime", "risk": "automated search state, retry, scoring, and checkpoint behavior", - "required_tests": ["tests/test_auto_obliterate.py", "tests/test_remaining_cpu_contracts.py"], - "conditional_gates": ["model-download-runtime"] + "required_tests": [ + "tests/test_auto_obliterate.py", + "tests/test_remaining_cpu_contracts.py" + ], + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/bayesian_optimizer.py", @@ -408,7 +500,9 @@ "tests/test_module_imports.py", "tests/conditional/test_model_download_runtime.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/evaluation/baselines.py", @@ -419,7 +513,9 @@ "tests/test_evaluation_reporting_contracts.py", "tests/conditional/test_model_download_runtime.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/evaluation/evaluator.py", @@ -430,7 +526,9 @@ "tests/test_evaluation_reporting_contracts.py", "tests/test_offline_integration.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/informed_pipeline.py", @@ -441,14 +539,21 @@ "tests/test_informed_pipeline_contracts.py", "tests/test_offline_integration.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/lora_ablation.py", "risk_class": "conditional-runtime", "risk": "optional adapter construction and validation against live projections", - "required_tests": ["tests/test_module_imports.py", "tests/conditional/test_model_download_runtime.py"], - "conditional_gates": ["model-download-runtime"] + "required_tests": [ + "tests/test_module_imports.py", + "tests/conditional/test_model_download_runtime.py" + ], + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/sweep.py", @@ -459,14 +564,21 @@ "tests/test_sweep_contracts.py", "tests/conditional/test_model_download_runtime.py" ], - "conditional_gates": ["model-download-runtime"] + "conditional_gates": [ + "model-download-runtime" + ] }, { "path": "obliteratus/evaluation/heretic_eval.py", "risk_class": "mixed-runtime", "risk": "publication evaluation prompts, classifiers, scoring, and external adapters", - "required_tests": ["tests/test_heretic_eval.py", "tests/conditional/test_external_evaluation_runtime.py"], - "conditional_gates": ["external-evaluation"] + "required_tests": [ + "tests/test_heretic_eval.py", + "tests/conditional/test_external_evaluation_runtime.py" + ], + "conditional_gates": [ + "external-evaluation" + ] }, { "path": "obliteratus/evaluation/lm_eval_integration.py", @@ -477,56 +589,93 @@ "tests/test_module_imports.py", "tests/conditional/test_external_evaluation_runtime.py" ], - "conditional_gates": ["external-evaluation"] + "conditional_gates": [ + "external-evaluation" + ] }, { "path": "obliteratus/tourney.py", "risk_class": "mixed-runtime", "risk": "multi-model tournament mutation, comparison, and optional publication", - "required_tests": ["tests/test_tourney_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_external_evaluation_runtime.py"], - "conditional_gates": ["external-evaluation"] + "required_tests": [ + "tests/test_tourney_contracts.py", + "tests/test_module_imports.py", + "tests/conditional/test_external_evaluation_runtime.py" + ], + "conditional_gates": [ + "external-evaluation" + ] }, { "path": "obliteratus/tourney_contracts.py", "risk_class": "cpu-contract", "risk": "checkpoint document parsing and version validation", - "required_tests": ["tests/test_tourney_contracts.py"], + "required_tests": [ + "tests/test_tourney_contracts.py" + ], "conditional_gates": [] }, { "path": "obliteratus/bestiary_sync.py", "risk_class": "conditional-runtime", "risk": "external catalog synchronization and malformed service responses", - "required_tests": ["tests/test_bestiary_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"], - "conditional_gates": ["network-services"] + "required_tests": [ + "tests/test_bestiary_contracts.py", + "tests/test_module_imports.py", + "tests/conditional/test_network_services.py" + ], + "conditional_gates": [ + "network-services" + ] }, { "path": "obliteratus/models_client.py", "risk_class": "conditional-runtime", "risk": "catalog resolution across operator files and network services", - "required_tests": ["tests/test_models_client_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"], - "conditional_gates": ["network-services"] + "required_tests": [ + "tests/test_models_client_contracts.py", + "tests/test_module_imports.py", + "tests/conditional/test_network_services.py" + ], + "conditional_gates": [ + "network-services" + ] }, { "path": "obliteratus/watchtower.py", "risk_class": "conditional-runtime", "risk": "scheduled scans, queue state, and live service responses", - "required_tests": ["tests/test_watchtower_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_network_services.py"], - "conditional_gates": ["network-services"] + "required_tests": [ + "tests/test_watchtower_contracts.py", + "tests/test_module_imports.py", + "tests/conditional/test_network_services.py" + ], + "conditional_gates": [ + "network-services" + ] }, { "path": "obliteratus/service_contracts.py", "risk_class": "cpu-contract", "risk": "catalog validation, external record normalization, and scheduler lifecycle decisions", - "required_tests": ["tests/test_models_client_contracts.py", "tests/test_watchtower_contracts.py"], + "required_tests": [ + "tests/test_models_client_contracts.py", + "tests/test_watchtower_contracts.py" + ], "conditional_gates": [] }, { "path": "obliteratus/interactive.py", "risk_class": "mixed-runtime", "risk": "interactive terminal prompts and operator decision flow", - "required_tests": ["tests/test_cli.py", "tests/test_interactive_contracts.py", "tests/conditional/test_operator_ui.py"], - "conditional_gates": ["operator-ui"] + "required_tests": [ + "tests/test_cli.py", + "tests/test_interactive_contracts.py", + "tests/conditional/test_operator_ui.py" + ], + "conditional_gates": [ + "operator-ui" + ] }, { "path": "obliteratus/local_ui.py", @@ -538,28 +687,54 @@ "tests/test_local_ui_portability.py", "tests/conditional/test_operator_ui.py" ], - "conditional_gates": ["operator-ui"] + "conditional_gates": [ + "operator-ui" + ] }, { "path": "obliteratus/ui_watchtower.py", "risk_class": "mixed-runtime", "risk": "service-backed UI tabs and scheduler controls", - "required_tests": ["tests/test_ui_watchtower_contracts.py", "tests/test_module_imports.py", "tests/conditional/test_operator_ui.py"], - "conditional_gates": ["operator-ui"] + "required_tests": [ + "tests/test_ui_watchtower_contracts.py", + "tests/test_module_imports.py", + "tests/conditional/test_operator_ui.py" + ], + "conditional_gates": [ + "operator-ui" + ] }, { "path": "obliteratus/remote.py", "risk_class": "mixed-runtime", "risk": "SSH discovery, quoting, execution, cancellation, and result synchronization", - "required_tests": ["tests/test_remote_boundaries.py", "tests/conditional/test_remote_runtime.py"], - "conditional_gates": ["remote-execution"] + "required_tests": [ + "tests/test_remote_boundaries.py", + "tests/conditional/test_remote_runtime.py" + ], + "conditional_gates": [ + "remote-execution" + ] }, { "path": "obliteratus/remote_contracts.py", "risk_class": "cpu-contract", "risk": "remote target validation, GPU selection normalization, and shell-safe command construction", - "required_tests": ["tests/test_remote_contracts.py", "tests/test_remote_boundaries.py"], + "required_tests": [ + "tests/test_remote_contracts.py", + "tests/test_remote_boundaries.py" + ], "conditional_gates": [] + }, + { + "path": "obliteratus/restore_multimodal.py", + "risk_class": "cpu-contract", + "risk": "merge abliterated text weights back into full multimodal model", + "required_tests": [ + "tests/test_module_imports.py" + ], + "conditional_gates": [], + "contract_owner": "OBLITERATUS maintainers" } ] } diff --git a/obliteratus/cli.py b/obliteratus/cli.py index c72fc88..64ca4cd 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -413,6 +413,18 @@ def main(argv: list[str] | None = None): "--dir", type=str, default="community_results", help="Directory containing contribution JSON files", ) + + # --- restore-multimodal --- + restore_parser = subparsers.add_parser( + "restore-multimodal", + help="Restore vision/MTP/linear-attn components after text-only abliteration", + ) + restore_parser.add_argument("--abliterated", required=True, + help="Text-only abliterated model directory") + restore_parser.add_argument("--stock", required=True, + help="Stock (full) model directory or HF repo") + restore_parser.add_argument("--output", required=True, + help="Output directory for merged model") aggregate_parser.add_argument( "--format", choices=["summary", "latex"], @@ -541,6 +553,12 @@ def main(argv: list[str] | None = None): _cmd_report(args) elif args.command == "aggregate": _cmd_aggregate(args) + elif args.command == "restore-multimodal": + from obliteratus.restore_multimodal import restore_multimodal + import logging + logging.basicConfig(level=logging.INFO, format="%(message)s") + result = restore_multimodal(args.abliterated, args.stock, args.output) + print(f"\nDone: {result['replaced']} abliterated + {result['kept']} stock = {result['total']} total") elif args.command == "ui": _cmd_ui(args) elif args.command == "recommend": diff --git a/obliteratus/restore_multimodal.py b/obliteratus/restore_multimodal.py new file mode 100644 index 0000000..ae9296e --- /dev/null +++ b/obliteratus/restore_multimodal.py @@ -0,0 +1,192 @@ +#!/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()