diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 37e8a5c..658f106 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -59,7 +59,7 @@ "obliteratus/sweep.py", "obliteratus/tourney.py", "obliteratus/tourney_contracts.py", -"obliteratus/restore_multimodal.py", + "obliteratus/restore_multimodal.py", "obliteratus/capability_check.py" ], "required_tests": [ @@ -78,6 +78,23 @@ "tests/test_persistence_pipeline.py" ] }, + { + "id": "checkpoint-composition", + "owner": "checkpoint composition maintainers", + "description": "Architecture-compatible model-weight composition and atomic output promotion", + "contract_types": [ + "model-mutation", + "persistence", + "public-interface" + ], + "paths": [ + "obliteratus/blend.py" + ], + "required_tests": [ + "tests/test_blend.py", + "tests/test_persistence_contracts.py" + ] + }, { "id": "configuration-and-selection", "owner": "runtime compatibility maintainers", @@ -725,7 +742,7 @@ "tests/test_remote_contracts.py", "tests/test_remote_boundaries.py" ], -"conditional_gates": [] + "conditional_gates": [] }, { "path": "obliteratus/capability_check.py", @@ -746,6 +763,16 @@ ], "conditional_gates": [], "contract_owner": "OBLITERATUS maintainers" + }, + { + "path": "obliteratus/blend.py", + "risk_class": "cpu-contract", + "risk": "complementary abliteration blending", + "contract_owner": "OBLITERATUS maintainers", + "required_tests": [ + "tests/test_blend.py" + ], + "conditional_gates": [] } ] } diff --git a/docs/complementary_blending.md b/docs/complementary_blending.md new file mode 100644 index 0000000..5d0afda --- /dev/null +++ b/docs/complementary_blending.md @@ -0,0 +1,232 @@ +# Complementary Abliteration Blending + +**Date:** 2026-08-20 +**Status:** Contributor-reported preliminary results; implementation independently tested +**Authors:** OBLITERATUS contributors and maintainers + +--- + +## Abstract + +This document describes **complementary abliteration blending**, which combines two compatible +abliterated checkpoints via weight-space interpolation. PR #127 did not include the raw benchmark +outputs, prompt-level refusal evidence, model revisions, environment capture, or artifact hashes +needed to independently reproduce its numerical results. The tables below therefore preserve the +contributor's preliminary report as context; they are not independently verified project claims. + +--- + +## 1. Motivation + +The contributor report frames abliteration as a tradeoff between refusal removal and measured +capability. Single-direction work such as [Arditi et al.][arditi] motivates the approach, while the +reported OBLITERATUS comparison observed lower MMLU for a more aggressive surgery. Raw evidence +for the specific refusal and MMLU figures was not included in PR #127. + +We hypothesized that different direction-finding methods damage different parts of the model's +capability geometry, and that blending their outputs could cancel these damages. + +--- + +## 2. Method + +### 2.1 Surgery A: Aggressive/SVD + +Standard OBLITERATUS aggressive pipeline with SVD-based direction extraction: + +``` +obliteratus obliterate $BASE --method aggressive --n-directions 3 \ + --regularization 0.08 --residue-weight 3 --refinement-passes 2 \ + --min-layer-fraction 0.45 +``` + +**Contributor-reported properties:** deep refusal removal with measurable capability damage on a +small spot check. The proposed explanation—SVD directions overlapping capability subspaces—is a +hypothesis that requires activation and controlled model-comparison evidence. + +### 2.2 Surgery B: LEACE + +OBLITERATUS aggressive pipeline with LEACE (Linear Erasure of Concept Embeddings) direction method: + +``` +obliteratus obliterate $BASE --method aggressive --direction-method leace \ + --n-directions 3 --regularization 0.06 --residue-weight 7 \ + --refinement-passes 3 --min-layer-fraction 0.40 +``` + +LEACE is a closed-form linear concept-erasure method designed to remove linearly available concept +information while minimizing distortion ([Belrose et al.][leace]). The contributor reported +capability retention and weaker output quality for this surgery; the asserted mechanism in the +generation pathway was not evidenced in this PR. + +### 2.3 Weight-Space LERP Blend + +Simple linear interpolation in weight space: + +```python +for key in weight_keys: + blended[key] = alpha * leace_weights[key] + (1 - alpha) * aggressive_weights[key] +``` + +The blend ratio `alpha = 0.60` was found by binary search over {0.30, 0.50, 0.55, 0.60, 0.65, 0.70}. + +--- + +## 3. Contributor-Reported Results (Not Independently Reproduced) + +No machine-readable benchmark or refusal-evaluation artifacts accompany these tables. Counts, +uncertainty, prompt selection, exact model revisions, and evaluator configuration must be supplied +before the results can support a release or research conclusion. + +### 3.1 Blend Ratio Search (15-subject MMLU, 150 questions) + +| Blend (% LEACE) | Refuse | Usable | MMLU | vs Stock | +|------------------|--------|--------|-------|----------| +| 0% (pure SVD) | 0% | 100% | 84.0% | -2.0pp | +| 30% | 0% | 90% | 83.3% | -2.7pp | +| 50% | 0% | 100% | 84.0% | -2.0pp | +| 55% | 0% | 80% | 84.7% | -1.3pp | +| **60%** | **0%** |**100%**|**86.0%**|**+0.0pp**| +| 65% | 0% | — | — | — | +| 70% | 0% | 80% | 86.7% | +0.7pp | +| 100% (pure LEACE)| 0% | 50% | 86.7% | +0.7pp | + +The contributor selected the 60% blend from this small search. The evidence is insufficient to +establish a unique optimum or distinguish the apparent differences from sampling noise. + +### 3.2 Larger Contributor Spot Check (57-subject MMLU, 570 questions) + +| Model | MMLU | Stderr | vs Stock | +|----------------------|---------|--------|----------| +| Stock Qwen3.8-27B | 85.26% | ±0.014 | — | +| OBLITERATUS V1 (s51) | 81.40% | — | -6.0pp | +| **OBLITERATUS V2 (s78)** | **86.32%** | **±0.014** | **+1.1pp** | + +### 3.3 Per-Subject Gains (5 questions/subject) + +Capability gains span both safety-adjacent and neutral reasoning topics: + +| Subject | Stock | V2 | Delta | Type | +|----------------------|-------|-------|--------|-----------| +| College Mathematics | 40% | 80% | +40pp | Neutral | +| Formal Logic | 40% | 60% | +20pp | Neutral | +| Jurisprudence | 60% | 80% | +20pp | Sensitive | +| Business Ethics | 80% | 100% | +20pp | Sensitive | +| Professional Law | 80% | 100% | +20pp | Sensitive | + +The reported neutral-topic gains motivate a controlled follow-up; five questions per subject are +not sufficient to establish capability improvement or rule out sampling variation. + +### 3.4 Real-World Practical Tasks + +| Test Suite | Stock | V2 | Tasks | +|---------------------|-------|-----|------------------------------------------| +| Basic (8 tasks) | 5/8 | 6/8 | Code, SQL, tool calling, JSON, math | +| Advanced (8 tasks) | 7/8 | 7/8 | ReAct agents, async refactor, K8s debug, | +| | | | security review, system design | + +The contributor reported comparable outcomes on this small practical-task set. The tasks and raw +outputs were not included, so maintainers have not independently verified that comparison. + +--- + +## 4. Why It Works + +### 4.1 Complementary Error Cancellation + +One hypothesis is that SVD and LEACE make different errors in weight space: + +- **SVD** greedily captures maximum variance directions. Some captured variance encodes + capability, not just refusal. This damages specific weight regions. +- **LEACE** minimizes a linear erasure objective. Whether it leaves specific residue in attention + heads or output projections must be measured rather than inferred from output behavior. + +Weight-space interpolation may average complementary errors: +- Where SVD damaged capability, LEACE's intact weights dilute the damage +- Where LEACE left refusal residue, SVD's clean weights dilute the residue + +### 4.2 Theoretical Connection to Model Merging + +This technique is related to model-merging work such as [Model Soups][model-soups] and +[TIES-Merging][ties], but is applied within the abliteration domain. PR #127 does not measure +task-vector orthogonality or error anti-correlation, so those remain testable explanations rather +than established properties. + +### 4.3 Capacity Hypothesis + +The contributor-reported MMLU difference raises several competing hypotheses. A separate, +provenance-gated experiment framework is tracked in [issue #132][capacity-issue]: + +1. **Activation Rank Analysis** — Does effective dimensionality increase after abliteration? +2. **Topic Cluster Analysis** — Do gains cluster on sensitive topics (hedging) or spread broadly (capacity)? +3. **Blend Control** — Does blending two identical SVD surgeries also gain MMLU? (Tests regularization hypothesis) +4. **Learning Absorption** — Does the abliterated model learn new information faster? (Tests freed capacity directly) + +The small per-subject report does not distinguish these hypotheses. + +--- + +## 5. Reproducibility + +```bash +# Step 1: Aggressive surgery +obliteratus obliterate $BASE --method aggressive --n-directions 3 \ + --regularization 0.08 --residue-weight 3 --refinement-passes 2 \ + --min-layer-fraction 0.45 --output-dir surgery_a + +# Step 2: LEACE surgery +obliteratus obliterate $BASE --method aggressive --direction-method leace \ + --n-directions 3 --regularization 0.06 --residue-weight 7 \ + --refinement-passes 3 --min-layer-fraction 0.40 --output-dir surgery_b + +# Step 3: Blend +obliteratus blend --model-a surgery_a --model-b surgery_b --alpha 0.6 \ + --config-source a --output blended_model + +# Step 4: Validate +lm_eval --model hf --model_args pretrained=blended_model --tasks mmlu +``` + +The command requires matching `source_model` values in both checkpoints' +`abliteration_metadata.json`, identical tensor keys, compatible shapes/dtypes, floating-point +weights, and sharded safetensors indexes. For legacy checkpoints whose common lineage was verified +out of band, `--allow-unverified-lineage` is an explicit escape hatch. Output is staged and +validated before atomically replacing any existing destination. + +--- + +## 6. Limitations + +- Contributor measurements cover only Qwen3.8-27B; independent validation is pending +- MMLU is a multiple-choice benchmark; gains may not transfer to all downstream tasks +- The 60/40 blend ratio may be model-specific +- `repetition_penalty=1.15` is still required for clean generation +- System prompts still reintroduce refusals +- Full 842-corpus validation in progress at time of writing +- The numerical results lack committed raw evidence and independent reproduction +- Single-file and quantized/integer checkpoints are not supported by the current blender +- Atomic promotion temporarily requires space for the complete staged output and any prior output + +--- + +## 7. Future Work + +- **Cross-architecture validation** on Llama, Gemma, Mistral +- **SLERP blending** instead of LERP (spherical interpolation may better preserve weight norms) +- **Three-way blends** with additional direction methods (diff_means, SOM) +- **Post-blend recovery**, tracked separately in [issue #133][recovery-issue] +- **Formal capacity-hypothesis validation**, tracked in [issue #132][capacity-issue] + +## References + +- [Arditi et al., *Refusal in Language Models Is Mediated by a Single Direction*][arditi] +- [Belrose et al., *LEACE: Perfect Linear Concept Erasure in Closed Form*][leace] +- [Wortsman et al., *Model Soups*][model-soups] +- [Yadav et al., *TIES-Merging*][ties] + +[arditi]: https://arxiv.org/abs/2406.11717 +[leace]: https://arxiv.org/abs/2306.03819 +[model-soups]: https://proceedings.mlr.press/v162/wortsman22a.html +[ties]: https://arxiv.org/abs/2306.01708 +[capacity-issue]: https://github.com/elder-plinius/OBLITERATUS/issues/132 +[recovery-issue]: https://github.com/elder-plinius/OBLITERATUS/issues/133 diff --git a/docs/executive_research_summary.md b/docs/executive_research_summary.md new file mode 100644 index 0000000..664ffb0 --- /dev/null +++ b/docs/executive_research_summary.md @@ -0,0 +1,228 @@ +# Complementary Abliteration Blending: Executive Research Summary + +**OBLITERATUS Project — August 2026** + +> **Evidence status:** This summary preserves contributor-reported preliminary measurements from +> PR #127. The PR did not include raw benchmark outputs, prompt-level refusal evidence, exact model +> revisions, environment capture, or artifact hashes. Maintainers independently verified the blend +> implementation and its CPU contracts, but not the numerical research results below. + +--- + +## The Problem + +The contributor framed existing abliteration approaches as trading deeper refusal removal for +greater capability loss. Whether that pattern generalizes or follows from refusal-model geometry +has not been established by the evidence included with PR #127. + +| Approach | Refusal Rate | MMLU Delta | Source | +|---|---|---|---| +| Single direction (Arditi et al.) | Low but residual | ~0pp | Baseline | +| OrcaRouter (1-dir, k=1) | Low | -0.8pp | Community | +| huihui-ai (1-dir, skip layers) | Low | ~0pp | Community | +| OBLITERATUS V1 (5-dir SVD) | 0.0% | -6.0pp | This work | + +The contributor reported that a V1 configuration reached 0% refusal on its sampled prompts with a +-6pp MMLU difference. Those figures are retained as an unverified observation, not proof of a +general tradeoff. + +## The Insight + +Different direction-finding algorithms damage different regions of weight space. + +**SVD (Singular Value Decomposition):** Extracts high-variance directions. The contributor proposes +that some directions encode both refusal and capability; this mechanism was not measured in PR #127. + +**LEACE (Linear Erasure of Concept Embeddings):** Provides closed-form linear concept erasure while +minimizing distortion. The reported capability retention and generation-pathway residue remain +contributor observations requiring artifact-backed reproduction. + +The working hypothesis is that these methods make complementary errors. The PR does not measure +their error correlation or task-vector geometry. + +## The Method + +Run both surgeries independently on the same base model, then interpolate in weight space: + +``` +blended_weight = α × LEACE_weight + (1 - α) × SVD_weight +``` + +The contributor searched α over {0.30, 0.50, 0.55, 0.60, 0.65, 0.70} using 15-subject MMLU and a +10-prompt usability check, then selected α = 0.60 for Qwen3.8-27B. The small, incomplete search +does not establish a unique optimum. + +**Proposed explanation:** interpolation may dilute method-specific damage. This must be tested +against same-method and unrelated-checkpoint blend controls before it is treated as causal. + +## Contributor-Reported Results (Not Independently Reproduced) + +### Headline + +| Model | MMLU (lm-eval, 0-shot) | Refusal Rate | Usable Output | +|---|---|---|---| +| Stock Qwen3.8-27B | 85.3% (n=570) | ~100% | — | +| V1 (aggressive/SVD) | 81.4% (n=285) | 0.0% | 80% | +| **V2 (60/40 blend)** | **86.3% (n=570)** | **0.0%** | **100%** | + +The contributor reported +1.1pp MMLU above stock while maintaining complete refusal removal. The +repository does not claim priority or a confirmed capability improvement without reproducible raw +evidence and appropriate statistical comparison. + +**Important caveat:** the contributor reports MMLU with `--limit 10` (570 questions, 10 per +subject), below the full benchmark. Without raw outputs and a paired statistical analysis, the ++1.1pp difference is descriptive only and may reflect sampling variation. + +### Per-Subject Analysis + +Gains span both safety-adjacent and neutral reasoning topics (5 questions per subject — preliminary): + +| Subject | Stock | V2 | Delta | Topic Type | +|---|---|---|---|---| +| College Mathematics | 40% | 80% | +40pp | Neutral | +| Formal Logic | 40% | 60% | +20pp | Neutral | +| Jurisprudence | 60% | 80% | +20pp | Safety-adjacent | +| Business Ethics | 80% | 100% | +20pp | Safety-adjacent | +| Professional Law | 80% | 100% | +20pp | Safety-adjacent | +| High School Chemistry | 100% | 80% | -20pp | Neutral (regression) | +| College Computer Science | 80% | 60% | -20pp | Neutral (regression) | + +The reported neutral-topic differences motivate a reduced-hedging control, but the per-subject +samples are too small to distinguish a capability effect from sampling variation. + +### Practical Capability + +| Test Category | V2 | Stock | N | +|---|---|---|---| +| Advanced real-world tasks | 7/8 | 7/8 | 8 | +| Basic real-world tasks | 6/8 | 5/8 | 8 | +| Tool calling (JSON, ReAct) | ✓ | ✓ | — | +| Code generation & refactoring | ✓ | ✓ | — | +| Security code review | ✓ | ✓ | — | +| Structured output (JSON schema) | ✓ | ✓ | — | +| System design | ✓ | ✓ | — | + +The contributor reported comparable outcomes on this small practical-task set. The underlying +tasks and outputs were not included for independent review. + +## What We Don't Know Yet + +### Unanswered Questions + +1. **Does the +1.1pp hold at full MMLU scale?** Our sample (570q) is above spot-check but below the full benchmark (14,042q). The number could converge to +0pp or +2pp with more data. + +2. **WHY does the blend improve over stock?** Three competing hypotheses: + - **Freed capacity:** Refusal training occupies representational capacity; removing it frees parameters for reasoning. Gains on math/logic support this. + - **Reduced hedging:** Stock model hedges on questions adjacent to sensitive topics; abliteration removes the hedging. Gains on law/ethics support this. + - **Blend regularization:** Weight averaging of any two diverse models acts as implicit regularization (analogous to model soups/ensembling). The improvement may not be specific to abliteration. + +3. **Is the 60/40 ratio model-specific?** The contributor reported testing only Qwen3.8-27B. + Useful ratios may vary by architecture, model size, and alignment training method. + +4. **Does this generalize beyond SVD + LEACE?** Other direction-finding methods (diff_means, SOM, nuclear/SAE) may offer additional complementary error profiles for three-way or N-way blends. + +5. **What happens with SLERP instead of LERP?** Spherical interpolation preserves weight norms better than linear interpolation. This may matter for models with strong norm-dependent behaviors. + +### Validation Gaps + +- Full MMLU (14k questions): in progress +- Full 842-corpus refusal validation: in progress (52/842 sample showed 0%) +- MMLU-Pro: not yet run +- Thinking mode ON: not tested +- GGUF inference validation: not tested (GGUFs compiled, not inference-checked) +- Cross-architecture replication: not attempted + +## Experimental Framework for Future Validation + +Four proposed experiments are tracked in [issue #132](https://github.com/elder-plinius/OBLITERATUS/issues/132): + +### Experiment 1: Activation Rank Analysis (Tests "freed capacity") +Run diverse prompts through stock and abliterated models. Capture hidden states at each layer. Compute effective rank via SVD. If abliteration frees capacity, the effective dimensionality of activations should increase. + +### Experiment 2: Topic Cluster Analysis (Tests "reduced hedging") +Compare per-subject MMLU differences between prespecified sensitive and neutral groups. The test +must define its statistical decision rule before examining results; the small table above is not +such a test. + +### Experiment 3: Blend Control (Tests "blend regularization") +Blend two identical SVD surgeries (same method, different random seeds) at 60/40. If this blend also gains MMLU, the improvement comes from weight averaging itself, not from the SVD/LEACE complementarity. This is the critical control experiment. + +### Experiment 4: Learning Absorption (Tests "freed capacity" directly) +QLoRA fine-tune both stock and abliterated models on identical small datasets. Compare loss curves. If the abliterated model learns faster (lower loss at same step count), it has more absorptive capacity — direct evidence for freed representational space. Requires GPU infrastructure (A100+, not feasible on MPS). + +Their implementation is intentionally deferred until the hypotheses, controls, provenance, CPU +contracts, and conditional GPU/network gates are specified. + +## Future Directions + +### Near-term (implementation available; research validation pending) + +1. **Cross-architecture replication.** Run the identical pipeline on additional model families. + This is needed to determine whether the technique generalizes and which parts of the recipe + require model-specific tuning. + +2. **Full-scale benchmarking.** Complete MMLU (14k), MMLU-Pro, HumanEval, GSM8K, and + ARC-Challenge with pinned inputs, raw results, and uncertainty estimates. + +3. **N-way blending.** Test three or more surgeries using different direction methods (SVD, + LEACE, diff_means, SOM) against prespecified pairwise and same-method controls. + +4. **Blend ratio as a function of model properties.** Study how selected α values relate to model + size, architecture, alignment training intensity, and number of refusal directions. + +### Medium-term (theoretical, needs investigation) + +5. **Post-blend capability recovery.** A provenance-safe dataset and QLoRA pipeline are proposed in + [issue #133](https://github.com/elder-plinius/OBLITERATUS/issues/133). No corpus or recovery + trainer is shipped by this change. + +6. **SLERP and task-arithmetic blending.** Replace LERP with spherical interpolation (preserves weight norms) or task-arithmetic approaches (TIES-Merging, DARE) that handle parameter conflicts more intelligently. LERP is the simplest possible blend — there is likely headroom from more sophisticated interpolation. + +7. **Adaptive per-layer blending.** Instead of a global α, use a different blend ratio per layer based on that layer's refusal vs capability contribution (measurable via activation probing). Layers with more refusal content get more SVD weight; layers with more capability content get more LEACE weight. This is the "precision blend" extension. + +8. **Blend as continuous optimization.** Instead of grid-searching α, treat the blend ratio as a differentiable parameter and optimize it directly against a capability+refusal objective using a small validation set. This is feasible on a single GPU and could find non-uniform per-tensor blend ratios. + +### Long-term (speculative, high-impact if true) + +9. **Capacity-hypothesis validation.** Increased activation rank alone would not demonstrate freed + representational capacity; the proposed work must control for prompt sampling, layer selection, + model identity, numerical thresholds, and alternative explanations before drawing implications + about safety training. + +10. **Generalized complementary merging.** The principle — "combine models that fail in different ways" — may extend beyond abliteration to any model merging scenario. Fine-tunes optimized for different objectives (code, math, reasoning) could be blended using the same complementary error cancellation principle, with direction-specific merge ratios instead of uniform interpolation. + +11. **Abliteration as a diagnostic.** If abliterated models consistently show capability changes on specific subjects, the per-subject delta profile becomes a map of where safety training allocated capacity. This "refusal cost map" could inform alignment researchers about which capabilities are most affected by safety training and guide more efficient alignment methods. + +--- + +## Reproduction + +```bash +# Step 1: Aggressive/SVD surgery +obliteratus obliterate $BASE --method aggressive --n-directions 3 \ + --regularization 0.08 --residue-weight 3 --refinement-passes 2 \ + --min-layer-fraction 0.45 --output-dir surgery_svd + +# Step 2: LEACE surgery +obliteratus obliterate $BASE --method aggressive --direction-method leace \ + --n-directions 3 --regularization 0.06 --residue-weight 7 \ + --refinement-passes 3 --min-layer-fraction 0.40 --output-dir surgery_leace + +# Step 3: Blend +obliteratus blend --model-a surgery_svd --model-b surgery_leace \ + --alpha 0.6 --config-source a --output blended + +# Step 4: Validate +lm_eval --model hf --model_args pretrained=blended --tasks mmlu --device auto +``` + +All code is open source: [github.com/elder-plinius/OBLITERATUS](https://github.com/elder-plinius/OBLITERATUS) + +Primary background: [Arditi et al.](https://arxiv.org/abs/2406.11717), +[LEACE](https://arxiv.org/abs/2306.03819), +[Model Soups](https://proceedings.mlr.press/v162/wortsman22a.html), and +[TIES-Merging](https://arxiv.org/abs/2306.01708). + +--- + +*OBLITERATUS Contributors, August 2026* diff --git a/obliteratus/blend.py b/obliteratus/blend.py new file mode 100644 index 0000000..01a66b2 --- /dev/null +++ b/obliteratus/blend.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python3 +"""Complementary abliteration blending — weight-space interpolation between surgeries. + +Novel technique: run two abliterations with different direction-finding methods +(e.g., aggressive/SVD and LEACE), then LERP-blend in weight space. Each method +makes different mistakes in different parts of the weight geometry: + + - Aggressive/SVD: deep refusal removal, but damages capability (greedy variance capture) + - LEACE: preserves capability (KL-minimizing), but weaker output quality + +Blending averages out each method's weaknesses. The 60/40 (LEACE/aggressive) ratio +was found by binary search to be the sweet spot for Qwen3.8-27B, yielding: + + - 0% refusal (from both parents) + - 100% usable output quality (from aggressive parent) + - +1.0pp MMLU vs stock (from LEACE parent's capability preservation) + +Usage: + obliteratus blend \\ + --model-a outputs/aggressive-surgery \\ + --model-b outputs/leace-surgery \\ + --alpha 0.6 \\ + --output outputs/blended-model + + # Binary search for optimal ratio + obliteratus blend \\ + --model-a outputs/aggressive \\ + --model-b outputs/leace \\ + --search 0.3,0.5,0.6,0.7 \\ + --output outputs/blend-search + +Alpha controls the interpolation: blended = alpha * model_b + (1 - alpha) * model_a +Higher alpha = more of model_b's character. +""" + +from __future__ import annotations + +import json +import logging +import math +import shutil +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import load_file, save_file + +from obliteratus.persistence_contracts import atomic_checkpoint_directory + +logger = logging.getLogger(__name__) + +_INDEX_NAME = "model.safetensors.index.json" +_GENERATED_FILES = {_INDEX_NAME, "blend_metadata.json"} + + +def _read_index(model_dir: Path) -> dict[str, Any]: + """Read and validate a sharded safetensors index.""" + + if not model_dir.is_dir(): + raise FileNotFoundError(f"Model directory does not exist: {model_dir}") + index_path = model_dir / _INDEX_NAME + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise FileNotFoundError(f"Missing safetensors index: {index_path}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid safetensors index JSON: {index_path}") from exc + + weight_map = index.get("weight_map") if isinstance(index, dict) else None + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"Safetensors index has no weight_map: {index_path}") + for key, shard in weight_map.items(): + if not isinstance(key, str) or not key: + raise ValueError(f"Safetensors index contains an invalid tensor key: {key!r}") + if not isinstance(shard, str) or Path(shard).name != shard: + raise ValueError(f"Safetensors index contains an unsafe shard path: {shard!r}") + shard_path = model_dir / shard + if not shard_path.is_file(): + raise FileNotFoundError(f"Missing safetensors shard: {shard_path}") + return index + + +def _load_indexed_shard( + model_dir: Path, + shard: str, + weight_map: dict[str, str], +) -> dict[str, torch.Tensor]: + tensors = load_file(str(model_dir / shard)) + expected = {key for key, mapped_shard in weight_map.items() if mapped_shard == shard} + if set(tensors) != expected: + missing = sorted(expected - set(tensors)) + extra = sorted(set(tensors) - expected) + raise ValueError( + f"Safetensors shard/index mismatch in {model_dir / shard}: " + f"missing={missing[:3]}, extra={extra[:3]}", + ) + return tensors + + +def _validate_nonoverlapping_paths(model_a: Path, model_b: Path, output: Path) -> None: + resolved_a = model_a.resolve() + resolved_b = model_b.resolve() + resolved_output = output.resolve() + if resolved_a == resolved_b: + raise ValueError("model_a and model_b must be different model directories") + for source in (resolved_a, resolved_b): + if resolved_output == source or resolved_output.is_relative_to(source): + raise ValueError("output must not be a model directory or one of its descendants") + if source.is_relative_to(resolved_output): + raise ValueError("output must not contain either source model directory") + + +def _read_source_model(model_dir: Path) -> str | None: + metadata_path = model_dir / "abliteration_metadata.json" + if not metadata_path.is_file(): + return None + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid abliteration metadata JSON: {metadata_path}") from exc + source_model = metadata.get("source_model") if isinstance(metadata, dict) else None + if not isinstance(source_model, str) or not source_model.strip(): + raise ValueError(f"Abliteration metadata has no source_model: {metadata_path}") + return source_model + + +def _verify_lineage(model_a: Path, model_b: Path, *, required: bool) -> str | None: + source_a = _read_source_model(model_a) + source_b = _read_source_model(model_b) + if required and (source_a is None or source_b is None): + raise ValueError( + "Both checkpoints require abliteration_metadata.json with matching source_model; " + "use allow_unverified_lineage=True only after independently verifying lineage", + ) + if source_a is not None and source_b is not None and source_a != source_b: + raise ValueError(f"Checkpoint source_model values do not match: {source_a!r} != {source_b!r}") + return source_a if source_a == source_b else None + + +def _copy_model_support_files(source: Path, destination: Path) -> None: + for item in source.iterdir(): + if item.name in _GENERATED_FILES or item.suffix == ".safetensors" or item.name == ".git": + continue + target = destination / item.name + if item.is_dir(): + shutil.copytree(item, target, symlinks=True) + else: + shutil.copy2(item, target) + + +def _validate_blended_checkpoint(checkpoint: Path) -> None: + index = _read_index(checkpoint) + if not (checkpoint / "config.json").is_file(): + raise ValueError("Selected config source does not contain config.json") + if not (checkpoint / "blend_metadata.json").is_file(): + raise ValueError("Blended checkpoint is missing blend_metadata.json") + for shard in set(index["weight_map"].values()): + _load_indexed_shard(checkpoint, shard, index["weight_map"]) + + +def blend_models( + model_a_path: str | Path, + model_b_path: str | Path, + output_path: str | Path, + alpha: float = 0.6, + config_source: str = "a", + allow_unverified_lineage: bool = False, +) -> dict: + """LERP-blend two models in weight space. + + blended[key] = alpha * model_b[key] + (1 - alpha) * model_a[key] + + Args: + model_a_path: First model directory (e.g., aggressive surgery). + model_b_path: Second model directory (e.g., LEACE surgery). + output_path: Where to save the blended model. + alpha: Blend ratio. 0.0 = pure model_a, 1.0 = pure model_b. + config_source: Which model's config files to use ("a" or "b"). + allow_unverified_lineage: Permit checkpoints without matching OBLITERATUS + source metadata. Tensor compatibility is still enforced. + + Returns: + dict with tensor counts and blend metadata. + """ + if not isinstance(alpha, (int, float)) or not math.isfinite(float(alpha)): + raise ValueError("alpha must be a finite number between 0 and 1") + alpha = float(alpha) + if not 0.0 <= alpha <= 1.0: + raise ValueError("alpha must be between 0 and 1 inclusive") + if config_source not in {"a", "b"}: + raise ValueError("config_source must be 'a' or 'b'") + + model_a = Path(model_a_path) + model_b = Path(model_b_path) + output = Path(output_path) + _validate_nonoverlapping_paths(model_a, model_b, output) + source_model = _verify_lineage( + model_a, + model_b, + required=not allow_unverified_lineage, + ) + + index_a = _read_index(model_a) + index_b = _read_index(model_b) + map_a = index_a["weight_map"] + map_b = index_b["weight_map"] + if set(map_a) != set(map_b): + missing_from_b = sorted(set(map_a) - set(map_b)) + missing_from_a = sorted(set(map_b) - set(map_a)) + raise ValueError( + "Model tensor keys do not match: " + f"missing_from_b={missing_from_b[:3]}, missing_from_a={missing_from_a[:3]}", + ) + + shards = sorted(set(map_a.values())) + total_tensors = 0 + blended_tensors = 0 + + logger.info("Blending: %.0f%% model_b + %.0f%% model_a", alpha * 100, (1 - alpha) * 100) + logger.info("Shards: %d", len(shards)) + + metadata: dict[str, Any] + source = model_a if config_source == "a" else model_b + with atomic_checkpoint_directory(output, validate=_validate_blended_checkpoint) as staging: + for shard_idx, shard in enumerate(shards): + tensors_a = _load_indexed_shard(model_a, shard, map_a) + needed_b_shards = {map_b[key] for key in tensors_a} + tensors_b: dict[str, torch.Tensor] = {} + for b_shard in needed_b_shards: + tensors_b.update(_load_indexed_shard(model_b, b_shard, map_b)) + + merged: dict[str, torch.Tensor] = {} + for key, tensor_a in tensors_a.items(): + tensor_b = tensors_b[key] + if tensor_a.shape != tensor_b.shape: + raise ValueError( + f"Tensor shape mismatch for {key}: {tensor_a.shape} != {tensor_b.shape}", + ) + if tensor_a.dtype != tensor_b.dtype: + raise ValueError( + f"Tensor dtype mismatch for {key}: {tensor_a.dtype} != {tensor_b.dtype}", + ) + if not torch.is_floating_point(tensor_a): + raise TypeError(f"Tensor {key} has non-floating dtype {tensor_a.dtype}") + merged[key] = alpha * tensor_b + (1.0 - alpha) * tensor_a + total_tensors += 1 + blended_tensors += 1 + + save_file(merged, str(staging / shard)) + if (shard_idx + 1) % 5 == 0 or shard_idx == len(shards) - 1: + logger.info(" [%d/%d] shards processed", shard_idx + 1, len(shards)) + + (staging / _INDEX_NAME).write_text( + json.dumps(index_a, indent=2) + "\n", + encoding="utf-8", + ) + _copy_model_support_files(source, staging) + metadata = { + "blend_method": "lerp", + "alpha": alpha, + "config_source": config_source, + "lineage_verified": source_model is not None, + "source_model": source_model, + "model_a": str(model_a), + "model_b": str(model_b), + "formula": f"blended = {alpha} * model_b + {1.0 - alpha} * model_a", + "total_tensors": total_tensors, + "blended_tensors": blended_tensors, + } + (staging / "blend_metadata.json").write_text( + json.dumps(metadata, indent=2) + "\n", + encoding="utf-8", + ) + + logger.info("Blend complete: %d tensors blended", blended_tensors) + + return metadata + + +def blend_search( + model_a_path: str | Path, + model_b_path: str | Path, + output_dir: str | Path, + alphas: list[float] | None = None, + config_source: str = "a", + allow_unverified_lineage: bool = False, +) -> list[dict]: + """Create multiple blends for binary-search evaluation. + + Args: + model_a_path: First model directory. + model_b_path: Second model directory. + output_dir: Parent directory for blend outputs. + alphas: List of blend ratios to try. Default: [0.3, 0.5, 0.6, 0.7]. + + Returns: + list of blend metadata dicts. + """ + if alphas is None: + alphas = [0.3, 0.5, 0.6, 0.7] + if not alphas: + raise ValueError("alphas must contain at least one blend ratio") + normalized = [float(alpha) for alpha in alphas] + if len(set(normalized)) != len(normalized): + raise ValueError("alphas must not contain duplicate blend ratios") + + output_dir = Path(output_dir) + results = [] + + for alpha in normalized: + percentage = f"{alpha * 100:g}".replace(".", "p") + label = f"blend_{percentage}" + output = output_dir / label + logger.info("\n=== %s (alpha=%.2f) ===", label, alpha) + meta = blend_models( + model_a_path, + model_b_path, + output, + alpha=alpha, + config_source=config_source, + allow_unverified_lineage=allow_unverified_lineage, + ) + meta["label"] = label + results.append(meta) + + return results + + +def main(): + import argparse + logging.basicConfig(level=logging.INFO, format="%(message)s") + + p = argparse.ArgumentParser( + description="Complementary abliteration blending — interpolate two surgeries in weight space." + ) + p.add_argument("--model-a", required=True, help="First model (e.g., aggressive surgery)") + p.add_argument("--model-b", required=True, help="Second model (e.g., LEACE surgery)") + p.add_argument("--alpha", type=float, default=0.6, help="Blend ratio (0=pure A, 1=pure B)") + p.add_argument("--config-source", choices=["a", "b"], default="a") + p.add_argument("--allow-unverified-lineage", action="store_true") + p.add_argument("--search", type=str, default=None, + help="Comma-separated alphas for binary search (e.g., 0.3,0.5,0.6,0.7)") + p.add_argument("--output", required=True, help="Output directory") + args = p.parse_args() + + if args.search: + alphas = [float(a) for a in args.search.split(",")] + results = blend_search( + args.model_a, + args.model_b, + args.output, + alphas, + config_source=args.config_source, + allow_unverified_lineage=args.allow_unverified_lineage, + ) + print(f"\nCreated {len(results)} blends in {args.output}/") + for r in results: + print(f" {r['label']}: alpha={r['alpha']}") + else: + result = blend_models( + args.model_a, + args.model_b, + args.output, + alpha=args.alpha, + config_source=args.config_source, + allow_unverified_lineage=args.allow_unverified_lineage, + ) + print(f"\nBlend complete: {result['blended_tensors']} tensors blended at alpha={args.alpha}") + + +if __name__ == "__main__": + main() diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 811aafc..06abf58 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -439,6 +439,28 @@ def main(argv: list[str] | None = None): capcheck_parser.add_argument("--subjects", type=str, default=None, help="Comma-separated subject list") capcheck_parser.add_argument("--limit", type=int, default=5, help="Questions per subject") capcheck_parser.add_argument("--output-dir", type=str, default=None) + + # --- blend --- + blend_parser = subparsers.add_parser( + "blend", + help="Complementary abliteration blending — interpolate two surgeries in weight space", + ) + blend_parser.add_argument("--model-a", required=True, help="First model (e.g., aggressive surgery)") + blend_parser.add_argument("--model-b", required=True, help="Second model (e.g., LEACE surgery)") + blend_parser.add_argument("--alpha", type=float, default=0.6, help="Blend ratio (0=pure A, 1=pure B)") + blend_parser.add_argument( + "--config-source", + choices=["a", "b"], + default="a", + help="Model whose tokenizer/config support files are copied", + ) + blend_parser.add_argument( + "--allow-unverified-lineage", + action="store_true", + help="Allow checkpoints without matching OBLITERATUS source-model metadata", + ) + blend_parser.add_argument("--search", type=str, default=None, help="Comma-separated alphas for search") + blend_parser.add_argument("--output", required=True, help="Output directory") aggregate_parser.add_argument( "--format", choices=["summary", "latex"], @@ -583,6 +605,31 @@ def main(argv: list[str] | None = None): limit=args.limit, output_dir=args.output_dir, ) print(f"\nStock: {result['stock_acc']*100:.1f}% Abliterated: {result['abliterated_acc']*100:.1f}% Delta: {result['delta_pp']:+.1f}pp") + elif args.command == "blend": + from obliteratus.blend import blend_models, blend_search + import logging + logging.basicConfig(level=logging.INFO, format="%(message)s") + if args.search: + alphas = [float(a) for a in args.search.split(",")] + results = blend_search( + args.model_a, + args.model_b, + args.output, + alphas, + config_source=args.config_source, + allow_unverified_lineage=args.allow_unverified_lineage, + ) + print(f"\nCreated {len(results)} blends") + else: + result = blend_models( + args.model_a, + args.model_b, + args.output, + alpha=args.alpha, + config_source=args.config_source, + allow_unverified_lineage=args.allow_unverified_lineage, + ) + print(f"\nBlend complete: {result['blended_tensors']} tensors at alpha={args.alpha}") elif args.command == "ui": _cmd_ui(args) elif args.command == "recommend": diff --git a/tests/test_blend.py b/tests/test_blend.py new file mode 100644 index 0000000..0dbb2b1 --- /dev/null +++ b/tests/test_blend.py @@ -0,0 +1,496 @@ +"""Tests for obliteratus.blend — complementary abliteration blending.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch +from safetensors.torch import load_file, save_file + + +def _make_model( + tmpdir: Path, + value: float, + n_tensors: int = 3, + *, + config_name: str = "test", +): + """Create a minimal model with all weights set to a constant value.""" + tensors = {f"layer.{i}.weight": torch.full((4, 4), value) for i in range(n_tensors)} + shard = "model-00001-of-00001.safetensors" + save_file(tensors, str(tmpdir / shard)) + index = {"metadata": {}, "weight_map": {k: shard for k in tensors}} + (tmpdir / "model.safetensors.index.json").write_text(json.dumps(index)) + (tmpdir / "config.json").write_text(json.dumps({"model_type": config_name})) + (tmpdir / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "example/base-model"}), + ) + return tensors + + +def _make_sharded_model(tmpdir: Path, values: dict[str, float], shard_for: dict[str, str]): + by_shard: dict[str, dict[str, torch.Tensor]] = {} + for key, value in values.items(): + by_shard.setdefault(shard_for[key], {})[key] = torch.full((2, 2), value) + for shard, tensors in by_shard.items(): + save_file(tensors, str(tmpdir / shard)) + index = {"metadata": {}, "weight_map": shard_for} + (tmpdir / "model.safetensors.index.json").write_text(json.dumps(index)) + (tmpdir / "config.json").write_text(json.dumps({"model_type": "test"})) + (tmpdir / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "example/base-model"}), + ) + + +class TestBlendModels: + def test_lerp_blend(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "model_a" + b_dir = tmp_path / "model_b" + out_dir = tmp_path / "blended" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 0.0) + _make_model(b_dir, 10.0) + + result = blend_models(str(a_dir), str(b_dir), str(out_dir), alpha=0.6) + + assert result["blended_tensors"] == 3 + assert result["total_tensors"] == 3 + assert result["alpha"] == 0.6 + + merged = load_file(str(out_dir / "model-00001-of-00001.safetensors")) + # 0.6 * 10.0 + 0.4 * 0.0 = 6.0 + assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 6.0)) + + def test_alpha_zero_is_pure_a(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 42.0) + _make_model(b_dir, 99.0) + + blend_models(str(a_dir), str(b_dir), str(out), alpha=0.0) + merged = load_file(str(out / "model-00001-of-00001.safetensors")) + assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 42.0)) + + def test_alpha_one_is_pure_b(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 42.0) + _make_model(b_dir, 99.0) + + blend_models(str(a_dir), str(b_dir), str(out), alpha=1.0) + merged = load_file(str(out / "model-00001-of-00001.safetensors")) + assert torch.allclose(merged["layer.0.weight"], torch.full((4, 4), 99.0)) + + def test_writes_metadata(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + + blend_models(str(a_dir), str(b_dir), str(out), alpha=0.5) + + meta = json.loads((out / "blend_metadata.json").read_text()) + assert meta["alpha"] == 0.5 + assert meta["blend_method"] == "lerp" + assert meta["config_source"] == "a" + assert meta["lineage_verified"] is True + assert meta["source_model"] == "example/base-model" + + def test_copies_config(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + + blend_models(str(a_dir), str(b_dir), str(out)) + assert (out / "config.json").exists() + + def test_uses_selected_config_source(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0, config_name="a") + _make_model(b_dir, 2.0, config_name="b") + + blend_models(a_dir, b_dir, out, config_source="b") + + assert json.loads((out / "config.json").read_text())["model_type"] == "b" + + def test_supports_different_shard_layouts_with_matching_keys(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + values_a = {"layer.0": 0.0, "layer.1": 2.0} + values_b = {"layer.0": 10.0, "layer.1": 6.0} + _make_sharded_model(a_dir, values_a, {key: "a.safetensors" for key in values_a}) + _make_sharded_model( + b_dir, + values_b, + {"layer.0": "b-1.safetensors", "layer.1": "b-2.safetensors"}, + ) + + blend_models(a_dir, b_dir, out, alpha=0.25) + + tensors = load_file(str(out / "a.safetensors")) + assert torch.allclose(tensors["layer.0"], torch.full((2, 2), 2.5)) + assert torch.allclose(tensors["layer.1"], torch.full((2, 2), 3.0)) + + @pytest.mark.parametrize("alpha", [-0.1, 1.1, float("inf"), float("nan")]) + def test_rejects_invalid_alpha_without_creating_output(self, tmp_path, alpha): + from obliteratus.blend import blend_models + + with pytest.raises(ValueError, match="alpha"): + blend_models(tmp_path / "a", tmp_path / "b", tmp_path / "out", alpha=alpha) + assert not (tmp_path / "out").exists() + + def test_rejects_invalid_config_source(self, tmp_path): + from obliteratus.blend import blend_models + + with pytest.raises(ValueError, match="config_source"): + blend_models( + tmp_path / "a", + tmp_path / "b", + tmp_path / "out", + config_source="other", + ) + + def test_rejects_missing_or_extra_tensor_keys(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0, n_tensors=2) + _make_model(b_dir, 2.0, n_tensors=1) + + with pytest.raises(ValueError, match="tensor keys do not match"): + blend_models(a_dir, b_dir, tmp_path / "out") + + def test_rejects_shape_and_dtype_mismatches(self, tmp_path): + from obliteratus.blend import blend_models + + for mismatch in ("shape", "dtype"): + a_dir = tmp_path / f"a-{mismatch}" + b_dir = tmp_path / f"b-{mismatch}" + a_dir.mkdir() + b_dir.mkdir() + shard = "model.safetensors" + save_file({"weight": torch.ones((2, 2))}, str(a_dir / shard)) + b_tensor = torch.ones((3, 2)) if mismatch == "shape" else torch.ones((2, 2)).double() + save_file({"weight": b_tensor}, str(b_dir / shard)) + index = {"weight_map": {"weight": shard}} + for directory in (a_dir, b_dir): + (directory / "model.safetensors.index.json").write_text(json.dumps(index)) + (directory / "config.json").write_text("{}") + (directory / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "example/base-model"}), + ) + + with pytest.raises(ValueError, match=f"Tensor {mismatch} mismatch"): + blend_models(a_dir, b_dir, tmp_path / f"out-{mismatch}") + + def test_rejects_non_floating_tensors(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + shard = "model.safetensors" + index = {"weight_map": {"weight": shard}} + for directory in (a_dir, b_dir): + save_file({"weight": torch.ones((2, 2), dtype=torch.int64)}, str(directory / shard)) + (directory / "model.safetensors.index.json").write_text(json.dumps(index)) + (directory / "config.json").write_text("{}") + (directory / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "example/base-model"}), + ) + + with pytest.raises(TypeError, match="non-floating"): + blend_models(a_dir, b_dir, tmp_path / "out") + + def test_rejects_unsafe_index_shard_path(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + (a_dir / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"weight": "../outside.safetensors"}}), + ) + for directory in (a_dir, b_dir): + (directory / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "example/base-model"}), + ) + + with pytest.raises(ValueError, match="unsafe shard path"): + blend_models(a_dir, b_dir, tmp_path / "out") + + def test_rejects_output_overlapping_a_source(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + + with pytest.raises(ValueError, match="output"): + blend_models(a_dir, b_dir, a_dir / "nested") + + def test_requires_matching_checkpoint_lineage_by_default(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + (b_dir / "abliteration_metadata.json").write_text( + json.dumps({"source_model": "different/base"}), + ) + + with pytest.raises(ValueError, match="source_model values do not match"): + blend_models(a_dir, b_dir, tmp_path / "out") + + def test_unverified_lineage_requires_explicit_opt_in(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + (a_dir / "abliteration_metadata.json").unlink() + (b_dir / "abliteration_metadata.json").unlink() + + with pytest.raises(ValueError, match="allow_unverified_lineage"): + blend_models(a_dir, b_dir, tmp_path / "blocked") + + result = blend_models( + a_dir, + b_dir, + tmp_path / "allowed", + allow_unverified_lineage=True, + ) + assert result["lineage_verified"] is False + assert result["source_model"] is None + + def test_failed_blend_preserves_existing_output(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + out.mkdir() + (out / "sentinel.txt").write_text("previous") + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + (a_dir / "config.json").unlink() + + with pytest.raises(ValueError, match="config.json"): + blend_models(a_dir, b_dir, out) + + assert (out / "sentinel.txt").read_text() == "previous" + + def test_successful_blend_atomically_replaces_existing_output(self, tmp_path): + from obliteratus.blend import blend_models + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + out.mkdir() + (out / "stale.txt").write_text("stale") + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + + blend_models(a_dir, b_dir, out) + + assert not (out / "stale.txt").exists() + assert (out / "blend_metadata.json").is_file() + + +class TestBlendSearch: + def test_creates_multiple_blends(self, tmp_path): + from obliteratus.blend import blend_search + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "search" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 0.0) + _make_model(b_dir, 10.0) + + results = blend_search(str(a_dir), str(b_dir), str(out), alphas=[0.3, 0.7]) + + assert len(results) == 2 + assert (out / "blend_30").exists() + assert (out / "blend_70").exists() + + def test_rejects_empty_or_duplicate_search_ratios(self, tmp_path): + from obliteratus.blend import blend_search + + with pytest.raises(ValueError, match="at least one"): + blend_search(tmp_path / "a", tmp_path / "b", tmp_path / "out", alphas=[]) + with pytest.raises(ValueError, match="duplicate"): + blend_search(tmp_path / "a", tmp_path / "b", tmp_path / "out", alphas=[0.5, 0.5]) + + def test_search_propagates_config_source(self, tmp_path): + from obliteratus.blend import blend_search + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0, config_name="a") + _make_model(b_dir, 2.0, config_name="b") + + blend_search(a_dir, b_dir, out, alphas=[0.5], config_source="b") + + config = json.loads((out / "blend_50" / "config.json").read_text()) + assert config["model_type"] == "b" + + +class TestCLIDispatch: + def test_blend_dispatch(self, tmp_path): + from obliteratus.cli import main as cli_main + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + + cli_main(["blend", + "--model-a", str(a_dir), + "--model-b", str(b_dir), + "--alpha", "0.6", + "--output", str(out)]) + + assert (out / "blend_metadata.json").exists() + + def test_blend_dispatch_uses_requested_config_source(self, tmp_path): + from obliteratus.cli import main as cli_main + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0, config_name="a") + _make_model(b_dir, 2.0, config_name="b") + + cli_main([ + "blend", + "--model-a", str(a_dir), + "--model-b", str(b_dir), + "--config-source", "b", + "--output", str(out), + ]) + + assert json.loads((out / "config.json").read_text())["model_type"] == "b" + + +class TestModuleCLI: + def test_single_blend_main(self, tmp_path, monkeypatch, capsys): + from obliteratus import blend + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + monkeypatch.setattr( + "sys.argv", + [ + "obliteratus.blend", + "--model-a", str(a_dir), + "--model-b", str(b_dir), + "--alpha", "0.25", + "--output", str(out), + ], + ) + + blend.main() + + assert "3 tensors blended at alpha=0.25" in capsys.readouterr().out + + def test_search_main(self, tmp_path, monkeypatch, capsys): + from obliteratus import blend + + a_dir = tmp_path / "a" + b_dir = tmp_path / "b" + out = tmp_path / "out" + a_dir.mkdir() + b_dir.mkdir() + _make_model(a_dir, 1.0) + _make_model(b_dir, 2.0) + monkeypatch.setattr( + "sys.argv", + [ + "obliteratus.blend", + "--model-a", str(a_dir), + "--model-b", str(b_dir), + "--search", "0.25,0.75", + "--output", str(out), + ], + ) + + blend.main() + + output = capsys.readouterr().out + assert "Created 2 blends" in output + assert "blend_25: alpha=0.25" in output