abliterax/ ├── pyproject.toml ├── README.md ├── abliterax/ │ ├── __init__.py │ ├── errors.py │ ├── logging_utils.py │ ├── config.py │ ├── model_loader.py │ ├── architecture_registry.py │ ├── dataset.py │ ├── activation_collector.py │ ├── geometry.py │ ├── interventions.py │ ├── benchmarks.py │ ├── optimizer.py │ ├── exporter.py │ └── cli.py ├── tests/ │ ├── test_geometry.py │ ├── test_interventions.py │ └── test_config.py ├── examples/ │ └── qwen3_0.6b_aggressive.toml └── notebooks/ └── abliterax_colab.ipynb 

[build-system] requires = ["hatchling"] build-backend = "hatchling.build" [project] name = "abliterax" version = "0.1.0" description = "Automated refusal-representation surgery for transformer LMs." requires-python = ">=3.10,<3.13" license = { text = "Apache-2.0" } dependencies = [ "torch>=2.2", "transformers>=4.51,<5", "accelerate>=0.33", "safetensors>=0.4", "optuna>=3.6", "numpy>=1.26", "scipy>=1.11", "datasets>=2.20", "huggingface-hub>=0.24", "sentencepiece>=0.2", "protobuf>=4.25", "tqdm>=4.66", ] [project.optional-dependencies] quant = ["bitsandbytes>=0.43"] eval = ["lm-eval>=0.4"] dev = ["pytest>=8"] [project.scripts] abliterax = "abliterax.cli:main" [tool.hatch.build.targets.wheel] packages = ["abliterax"] 

class AbliteraxError(Exception): """Base class for all pipeline errors.""" class ConfigError(AbliteraxError): """Raised when a configuration fails validation.""" class IncompatibleArchitecture(AbliteraxError): """Raised when the model architecture is not safely supported.""" class InsufficientResources(AbliteraxError): """Raised when VRAM/CPU estimates fall short of the plan.""" class VerificationFailed(AbliteraxError): """Raised when a post-intervention integrity check fails.""" 

import logging import sys def setup_logging(level: int = logging.INFO) -> None: logging.basicConfig( level=level, stream=sys.stderr, format="%(asctime)s %(levelname)-7s %(name)s :: %(message)s", datefmt="%H:%M:%S", ) def get_logger(name: str) -> logging.Logger: return logging.getLogger(name) 

from __future__ import annotations import dataclasses import tomllib from dataclasses import dataclass, field from pathlib import Path from typing import Literal from .errors import ConfigError DirectionMethod = Literal["diff_means", "svd", "whitened_svd", "fisher_lda"] Component = Literal["attn", "mlp"] RowNormMode = Literal["none", "pre", "full"] SaveMode = Literal["merged", "lora"] @dataclass class ModelSpec: model_id: str = "" revision: str = "main" dtype: Literal["float32", "bfloat16", "float16"] = "bfloat16" quantization: Literal["bnb_4bit"] | None = None trust_remote_code: bool = False @dataclass class ProbeConfig: good_dataset: str = "mlabonne/harmless_alpaca" good_column: str = "text" good_total: int = 512 good_probe: int = 256 bad_dataset: str = "mlabonne/harmful_behaviors" bad_column: str = "text" bad_total: int = 512 bad_probe: int = 256 max_seq_len: int = 512 use_chat_template: bool = True seed: int = 0 @dataclass class GeometryConfig: method: DirectionMethod = "whitened_svd" n_directions: int = 4 orthogonalize_good: bool = True winsorize: bool = True winsorize_quantile: float = 0.995 whiten_eps: float = 1e-4 random_svd_oversample: int = 16 random_svd_power_iters: int = 2 layer_stride: int = 1 @dataclass class InterventionConfig: row_norm_mode: RowNormMode = "full" project_bias: bool = True direction_index_mode: Literal["shared", "per_layer"] = "per_layer" max_norm_ratio: float = 1.10 verify_max_abs_delta: float = 1e-3 @dataclass class OptimizerConfig: n_trials: int = 16 n_startup_trials: int = 6 seed: int = 0 sampler: Literal["tpe", "random"] = "tpe" eval_batch_size: int = 4 eval_n_good: int = 64 eval_n_bad: int = 64 max_new_tokens: int = 64 max_weight_upper: float = 1.2 early_stop_patience: int = 8 post_search_top_k: int = 3 @dataclass class BenchmarkConfig: run_lm_eval: bool = True lm_eval_tasks: list[str] = field(default_factory=lambda: ["mmlu", "gsm8k", "humaneval"]) lm_eval_batch_size: int = 8 ppl_dataset: str = "wikitext" ppl_config: str = "wikitext-2-raw-v1" ppl_split: str = "test" ppl_n: int = 64 mmlu_min: float | None = None gsm8k_min: float | None = None humaneval_min: float | None = None ppl_max_ratio: float | None = 1.35 @dataclass class ExportConfig: out_dir: str = "./abliterax_out" save_mode: SaveMode = "merged" verify_reload: bool = True hub_repo: str | None = None hub_private: bool = True @dataclass class HardwareConfig: vram_gb: float | None = None allow_cpu: bool = False offload_dir: str | None = None reserve_gb: float = 3.0 @dataclass class Config: model: ModelSpec probe: ProbeConfig = field(default_factory=ProbeConfig) geometry: GeometryConfig = field(default_factory=GeometryConfig) intervention: InterventionConfig = field(default_factory=InterventionConfig) optimizer: OptimizerConfig = field(default_factory=OptimizerConfig) benchmarks: BenchmarkConfig = field(default_factory=BenchmarkConfig) export: ExportConfig = field(default_factory=ExportConfig) hardware: HardwareConfig = field(default_factory=HardwareConfig) def validate(self) -> None: if not self.model.model_id: raise ConfigError("model.model_id is required") if self.geometry.method not in ("diff_means", "svd", "whitened_svd", "fisher_lda"): raise ConfigError(f"unknown geometry.method={self.geometry.method!r}") if self.geometry.n_directions < 1: raise ConfigError("geometry.n_directions must be >= 1") if self.probe.good_probe < 8 or self.probe.bad_probe < 8: raise ConfigError("need at least 8 probe prompts per class") if self.probe.good_probe + self.probe.good_total > 0 and self.probe.good_probe >= self.probe.good_total: raise ConfigError("probe.good_probe must be < probe.good_total") if self.probe.bad_probe >= self.probe.bad_total: raise ConfigError("probe.bad_probe must be < probe.bad_total") if self.optimizer.max_weight_upper < 0: raise ConfigError("optimizer.max_weight_upper must be >= 0") if self.intervention.row_norm_mode not in ("none", "pre", "full"): raise ConfigError("intervention.row_norm_mode invalid") def to_dict(self) -> dict: return dataclasses.asdict(self) def load_config(path: str | Path) -> Config: data = tomllib.loads(Path(path).read_text(encoding="utf-8")) cfg = _build(dict(data)) cfg.validate() return cfg def _build(data: dict) -> Config: def take(key: str, cls): raw = data.pop(key, {}) or {} known = {f.name for f in dataclasses.fields(cls)} return cls(**{k: v for k, v in raw.items() if k in known}) model = take("model", ModelSpec) cfg = Config(model=model) cfg.probe = take("probe", ProbeConfig) cfg.geometry = take("geometry", GeometryConfig) cfg.intervention = take("intervention", InterventionConfig) cfg.optimizer = take("optimizer", OptimizerConfig) cfg.benchmarks = take("benchmarks", BenchmarkConfig) cfg.export = take("export", ExportConfig) cfg.hardware = take("hardware", HardwareConfig) if data: raise ConfigError(f"unknown top-level sections: {sorted(data)}") return cfg 

"""Architecture registry. Each resolver returns an `ArchitectureSpec` whose `get_*` callables fetch the actual `torch.nn.Linear` modules for a layer index. We only know the tensor mapping after iterating real blocks, so every resolver is validated by `verify()` before the pipeline proceeds. """ from __future__ import annotations from dataclasses import dataclass from typing import Callable import torch from torch import nn from .errors import IncompatibleArchitecture, VerificationFailed from .logging_utils import get_logger log = get_logger(__name__) @dataclass class ArchitectureSpec: model_type: str hidden_size: int n_layers: int is_moe: bool is_multimodal: bool get_block: Callable[[int], nn.Module] get_attn_out: Callable[[nn.Module], nn.Linear] get_mlp_down: Callable[[nn.Module], nn.Linear] note: str = "" def _text_conf(conf) -> dict: if hasattr(conf, "text_config") and conf.text_config is not None: return dict(conf.text_config.to_dict()) return dict(conf.to_dict()) def _check_linear(mod, hidden: int, where: str) -> nn.Linear: if not isinstance(mod, nn.Linear): raise VerificationFailed(f"{where}: expected nn.Linear, got {type(mod).__name__}") if mod.weight.ndim != 2: raise VerificationFailed(f"{where}: weight must be 2-D, got {tuple(mod.weight.shape)}") if mod.weight.shape[0] != hidden: raise VerificationFailed( f"{where}: weight out-dim {mod.weight.shape[0]} != hidden_size {hidden}" ) return mod def _llama_like(model, conf): c = _text_conf(conf) layers = model.model.layers return ArchitectureSpec( model_type=c["model_type"], hidden_size=c["hidden_size"], n_layers=len(layers), is_moe=False, is_multimodal=False, get_block=lambda i: layers[i], get_attn_out=lambda b: b.self_attn.o_proj, get_mlp_down=lambda b: b.mlp.down_proj, ) def _mixtral_like(model, conf<...