mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-07-11 22:56:32 +02:00
pliny
43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
"""Base class for ablation strategies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import abc
|
|
from dataclasses import dataclass
|
|
from typing import Any, Iterator
|
|
|
|
from obliteratus.models.loader import ModelHandle
|
|
|
|
|
|
@dataclass
|
|
class AblationSpec:
|
|
"""Describes one atomic ablation operation."""
|
|
|
|
strategy_name: str
|
|
component: str # human-readable name, e.g. "layer_3", "head_2_5"
|
|
description: str
|
|
metadata: dict[str, Any] | None = None
|
|
|
|
|
|
class AblationStrategy(abc.ABC):
|
|
"""Base class that all ablation strategies must implement."""
|
|
|
|
name: str = "base"
|
|
|
|
@abc.abstractmethod
|
|
def enumerate(self, handle: ModelHandle, **kwargs) -> list[AblationSpec]:
|
|
"""Return every possible ablation this strategy can perform on the model."""
|
|
|
|
@abc.abstractmethod
|
|
def apply(self, handle: ModelHandle, spec: AblationSpec) -> None:
|
|
"""Apply a single ablation in-place. The caller is responsible for
|
|
calling handle.restore() afterwards to undo the modification."""
|
|
|
|
def iterate(self, handle: ModelHandle, **kwargs) -> Iterator[AblationSpec]:
|
|
"""Convenience: yield specs one at a time, applying + restoring around each."""
|
|
specs = self.enumerate(handle, **kwargs)
|
|
for spec in specs:
|
|
self.apply(handle, spec)
|
|
yield spec
|
|
handle.restore()
|