Core

The core module defines the contracts that all methods share.

Most contributors should start with these classes:

• ModelWrapper: abstraction for model loading, hooks, cached inference, and generation.
• BaseProbe: interface for endogenous probes.
• BaseMonitor: interface for runtime safety monitors.
• BaseAttributor: interface for input or training-data attribution.
• SafetyReport: standard report format consumed by adapters such as FlagSafe.
• PipelineConfig: validated YAML configuration model.

Minimal BaseProbe:

from collections.abc import Sequence
from typing import Any

from SafeLens.core.base import BaseProbe, Batch, ModelWrapper, ProbeResult


class AlwaysSafeProbe(BaseProbe):
    def attach(self, model: ModelWrapper, layers: Sequence[int]) -> None:
        self.layers = list(layers)

    def detect(self, batch: Batch) -> ProbeResult:
        return ProbeResult(risk_score=0.0, critical_layers=self.layers)

    def intervene(self, batch: Batch, direction: Any, scale: float) -> None:
        pass

    def detach(self) -> None:
        self.layers = []

Minimal BaseMonitor:

from typing import Any

from SafeLens.core.base import BaseMonitor, Batch, ModelWrapper, MonitoringSignal, SafetyReport


class AlwaysSafeMonitor(BaseMonitor):
    def start_monitoring(self, model: ModelWrapper) -> None:
        self.signals: list[MonitoringSignal] = []

    def step(self, batch: Batch, model_output: Any = None) -> MonitoringSignal:
        signal = MonitoringSignal(name="always_safe", risk_score=0.0)
        self.signals.append(signal)
        return signal

    def report(self) -> SafetyReport:
        return SafetyReport(monitoring_signals=self.signals)

Minimal BaseAttributor:

from typing import Any

from SafeLens.core.base import AttributionResult, BaseAttributor, Batch


class EmptyAttributor(BaseAttributor):
    def attribute_training(self, batch: Batch, model_output: Any = None) -> AttributionResult:
        return AttributionResult(method="empty", attribution_score=0.0)

    def attribute_input(self, batch: Batch, model_output: Any = None) -> AttributionResult:
        return AttributionResult(method="empty", attribution_score=0.0)

Shared contracts for SafeLens methods, reports, and pipelines.

AttributionResult

Bases: SerializableModel

Attribution output for input or training-data influence.

Source code in src/SafeLens/core/base.py

class AttributionResult(SerializableModel):
    """Attribution output for input or training-data influence."""

    method: str
    attribution_score: float = Field(ge=0.0, le=1.0)
    tokens: list[TokenAttribution] = Field(default_factory=list)
    details: dict[str, Any] = Field(default_factory=dict)

BaseAttributor

Bases: ABC

Base class for input and training-data attribution methods.

Source code in src/SafeLens/core/base.py

class BaseAttributor(ABC):
    """Base class for input and training-data attribution methods."""

    name: ClassVar[str] = ""

    def __init__(self, config: Mapping[str, Any] | None = None) -> None:
        self.config = dict(config or {})

    def attach(self, model: ModelWrapper) -> None:
        """Attach the attributor to a loaded model wrapper."""
        _ = model

    def detach(self) -> None:
        """Release model references or runtime state held by the attributor."""
        return None

    @abstractmethod
    def attribute_training(self, batch: Batch, model_output: Any = None) -> AttributionResult:
        """Estimate influential training examples or sources for the batch."""

    @abstractmethod
    def attribute_input(self, batch: Batch, model_output: Any = None) -> AttributionResult:
        """Estimate input-token contribution to the observed risk."""

attach(model)

Attach the attributor to a loaded model wrapper.

Source code in src/SafeLens/core/base.py

def attach(self, model: ModelWrapper) -> None:
    """Attach the attributor to a loaded model wrapper."""
    _ = model

attribute_input(batch, model_output=None) abstractmethod

Estimate input-token contribution to the observed risk.

Source code in src/SafeLens/core/base.py

@abstractmethod
def attribute_input(self, batch: Batch, model_output: Any = None) -> AttributionResult:
    """Estimate input-token contribution to the observed risk."""

attribute_training(batch, model_output=None) abstractmethod

Estimate influential training examples or sources for the batch.

Source code in src/SafeLens/core/base.py

@abstractmethod
def attribute_training(self, batch: Batch, model_output: Any = None) -> AttributionResult:
    """Estimate influential training examples or sources for the batch."""

detach()

Release model references or runtime state held by the attributor.

Source code in src/SafeLens/core/base.py

def detach(self) -> None:
    """Release model references or runtime state held by the attributor."""
    return None

BaseMethodConfig

Bases: SerializableModel

Base configuration class for pluggable methods.

Source code in src/SafeLens/core/base.py

class BaseMethodConfig(SerializableModel):
    """Base configuration class for pluggable methods."""

    enabled: bool = True
    params: dict[str, Any] = Field(default_factory=dict)

BaseMonitor

Bases: ABC

Base class for generation-time safety monitors.

Source code in src/SafeLens/core/base.py

class BaseMonitor(ABC):
    """Base class for generation-time safety monitors."""

    name: ClassVar[str] = ""

    def __init__(self, config: Mapping[str, Any] | None = None) -> None:
        self.config = dict(config or {})

    @abstractmethod
    def start_monitoring(self, model: ModelWrapper) -> None:
        """Initialize monitor state for a model run."""

    @abstractmethod
    def step(self, batch: Batch, model_output: Any = None) -> MonitoringSignal:
        """Inspect one batch or generation step and emit a safety signal."""

    @abstractmethod
    def report(self) -> SafetyReport:
        """Return the monitor's aggregate safety report."""

report() abstractmethod

Return the monitor's aggregate safety report.

Source code in src/SafeLens/core/base.py

@abstractmethod
def report(self) -> SafetyReport:
    """Return the monitor's aggregate safety report."""

start_monitoring(model) abstractmethod

Initialize monitor state for a model run.

Source code in src/SafeLens/core/base.py

@abstractmethod
def start_monitoring(self, model: ModelWrapper) -> None:
    """Initialize monitor state for a model run."""

step(batch, model_output=None) abstractmethod

Inspect one batch or generation step and emit a safety signal.

Source code in src/SafeLens/core/base.py

@abstractmethod
def step(self, batch: Batch, model_output: Any = None) -> MonitoringSignal:
    """Inspect one batch or generation step and emit a safety signal."""

BaseProbe

Bases: ABC

Base class for endogenous safety probes.

Source code in src/SafeLens/core/base.py

class BaseProbe(ABC):
    """Base class for endogenous safety probes."""

    name: ClassVar[str] = ""

    def __init__(self, config: Mapping[str, Any] | None = None) -> None:
        self.config = dict(config or {})

    @abstractmethod
    def attach(self, model: ModelWrapper, layers: Sequence[LayerRef]) -> None:
        """Register hooks on target layers."""

    @abstractmethod
    def detect(self, batch: Batch) -> ProbeResult:
        """Compute safety risk from the current batch and cached state."""

    @abstractmethod
    def intervene(self, batch: Batch, direction: Any, scale: float) -> None:
        """Apply an activation-space intervention."""

    @abstractmethod
    def detach(self) -> None:
        """Remove probe hooks and clear runtime state."""

attach(model, layers) abstractmethod

Register hooks on target layers.

Source code in src/SafeLens/core/base.py

@abstractmethod
def attach(self, model: ModelWrapper, layers: Sequence[LayerRef]) -> None:
    """Register hooks on target layers."""

detach() abstractmethod

Remove probe hooks and clear runtime state.

Source code in src/SafeLens/core/base.py

@abstractmethod
def detach(self) -> None:
    """Remove probe hooks and clear runtime state."""

detect(batch) abstractmethod

Compute safety risk from the current batch and cached state.

Source code in src/SafeLens/core/base.py

@abstractmethod
def detect(self, batch: Batch) -> ProbeResult:
    """Compute safety risk from the current batch and cached state."""

intervene(batch, direction, scale) abstractmethod

Apply an activation-space intervention.

Source code in src/SafeLens/core/base.py

@abstractmethod
def intervene(self, batch: Batch, direction: Any, scale: float) -> None:
    """Apply an activation-space intervention."""

MethodSpec

Bases: SerializableModel

Name and config payload for a registered method.

Source code in src/SafeLens/core/base.py

class MethodSpec(SerializableModel):
    """Name and config payload for a registered method."""

    name: str = Field(min_length=1)
    config: dict[str, Any] = Field(default_factory=dict)

ModelLoadConfig

Bases: SerializableModel

Configuration for model wrapper construction.

Source code in src/SafeLens/core/base.py

class ModelLoadConfig(SerializableModel):
    """Configuration for model wrapper construction."""

    name: str = Field(default="dummy", min_length=1)
    source: str = Field(
        default="huggingface",
        description="Model backend selector.",
        json_schema_extra={"enum": list(SUPPORTED_MODEL_SOURCES)},
    )
    dtype: str = "float32"
    device: str | None = None
    revision: str | None = None
    cache_dir: str | None = None
    local_dir: str | None = None
    trust_remote_code: bool = False
    load_kwargs: dict[str, Any] = Field(default_factory=dict)
    tokenizer_kwargs: dict[str, Any] = Field(default_factory=dict)
    modelscope_kwargs: dict[str, Any] = Field(default_factory=dict)

    @field_validator("source")
    @classmethod
    def validate_source(cls, value: str) -> str:
        normalized = value.strip().lower()
        if normalized in SUPPORTED_MODEL_SOURCES:
            return normalized
        suggestion = get_close_matches(normalized, SUPPORTED_MODEL_SOURCES, n=1)
        hint = f" Did you mean {suggestion[0]!r}?" if suggestion else ""
        expected = ", ".join(SUPPORTED_MODEL_SOURCES)
        raise ValueError(f"Unsupported model.source {value!r}. Expected one of: {expected}.{hint}")

ModelWrapper

Bases: ABC

Abstract model interface used by probes, monitors, and attributors.

Source code in src/SafeLens/core/base.py

class ModelWrapper(ABC):
    """Abstract model interface used by probes, monitors, and attributors."""

    @abstractmethod
    def load_model(self) -> Any:
        """Load and return the underlying model object."""

    @abstractmethod
    def add_hook(
        self,
        layer: LayerRef,
        hook_fn: HookFn | None = None,
        *,
        hook: HookFn | None = None,
        dir: str = "fwd",
        is_permanent: bool = False,
        level: int | None = None,
        prepend: bool = False,
    ) -> Any:
        """Register a forward hook on a target layer."""

    def add_perma_hook(self, layer: LayerRef, hook_fn: HookFn) -> Any:
        """Register a persistent forward hook when the wrapper supports it."""
        return self.add_hook(layer, hook_fn, is_permanent=True)

    @abstractmethod
    def run_with_cache(
        self,
        batch: Batch,
        layers: Sequence[LayerRef] | None = None,
        **kwargs: Any,
    ) -> tuple[Any, Any]:
        """Run inference and optionally return cached activations for selected layers."""

    @abstractmethod
    def generate(self, prompt: str, **generation_kwargs: Any) -> Any:
        """Generate text or model outputs from a prompt."""

    @abstractmethod
    def remove_hooks(self) -> None:
        """Remove all active hooks managed by this wrapper."""

    def reset_hooks(
        self,
        *,
        clear_contexts: bool = True,
        direction: Any = None,
        dir: Any = None,
        including_permanent: bool = False,
        level: Any = None,
    ) -> None:
        """TransformerLens-compatible alias for clearing wrapper-managed hooks."""
        _ = clear_contexts, direction, dir, including_permanent, level
        self.remove_hooks()

add_hook(layer, hook_fn=None, *, hook=None, dir='fwd', is_permanent=False, level=None, prepend=False) abstractmethod

Register a forward hook on a target layer.

Source code in src/SafeLens/core/base.py

@abstractmethod
def add_hook(
    self,
    layer: LayerRef,
    hook_fn: HookFn | None = None,
    *,
    hook: HookFn | None = None,
    dir: str = "fwd",
    is_permanent: bool = False,
    level: int | None = None,
    prepend: bool = False,
) -> Any:
    """Register a forward hook on a target layer."""

add_perma_hook(layer, hook_fn)

Register a persistent forward hook when the wrapper supports it.

Source code in src/SafeLens/core/base.py

def add_perma_hook(self, layer: LayerRef, hook_fn: HookFn) -> Any:
    """Register a persistent forward hook when the wrapper supports it."""
    return self.add_hook(layer, hook_fn, is_permanent=True)

generate(prompt, **generation_kwargs) abstractmethod

Generate text or model outputs from a prompt.

Source code in src/SafeLens/core/base.py

@abstractmethod
def generate(self, prompt: str, **generation_kwargs: Any) -> Any:
    """Generate text or model outputs from a prompt."""

load_model() abstractmethod

Load and return the underlying model object.

Source code in src/SafeLens/core/base.py

@abstractmethod
def load_model(self) -> Any:
    """Load and return the underlying model object."""

remove_hooks() abstractmethod

Remove all active hooks managed by this wrapper.

Source code in src/SafeLens/core/base.py

@abstractmethod
def remove_hooks(self) -> None:
    """Remove all active hooks managed by this wrapper."""

reset_hooks(*, clear_contexts=True, direction=None, dir=None, including_permanent=False, level=None)

TransformerLens-compatible alias for clearing wrapper-managed hooks.

Source code in src/SafeLens/core/base.py

def reset_hooks(
    self,
    *,
    clear_contexts: bool = True,
    direction: Any = None,
    dir: Any = None,
    including_permanent: bool = False,
    level: Any = None,
) -> None:
    """TransformerLens-compatible alias for clearing wrapper-managed hooks."""
    _ = clear_contexts, direction, dir, including_permanent, level
    self.remove_hooks()

run_with_cache(batch, layers=None, **kwargs) abstractmethod

Run inference and optionally return cached activations for selected layers.

Source code in src/SafeLens/core/base.py

@abstractmethod
def run_with_cache(
    self,
    batch: Batch,
    layers: Sequence[LayerRef] | None = None,
    **kwargs: Any,
) -> tuple[Any, Any]:
    """Run inference and optionally return cached activations for selected layers."""

MonitoringSignal

Bases: SerializableModel

Per-step safety signal emitted by a monitor.

Source code in src/SafeLens/core/base.py

class MonitoringSignal(SerializableModel):
    """Per-step safety signal emitted by a monitor."""

    name: str
    risk_score: float = Field(ge=0.0, le=1.0)
    triggered: bool = False
    risk_category: list[str] = Field(default_factory=list)
    evidence_tokens: list[int] = Field(default_factory=list)
    details: dict[str, Any] = Field(default_factory=dict)

OutputConfig

Bases: SerializableModel

Output configuration for generated reports.

Source code in src/SafeLens/core/base.py

class OutputConfig(SerializableModel):
    """Output configuration for generated reports."""

    report_path: str = "./safety_scan.json"

PipelineConfig

Bases: SerializableModel

Top-level YAML config for safelens run.

Source code in src/SafeLens/core/base.py

class PipelineConfig(SerializableModel):
    """Top-level YAML config for `safelens run`."""

    model: ModelLoadConfig = Field(default_factory=ModelLoadConfig)
    pipeline: PipelineSectionConfig = Field(default_factory=PipelineSectionConfig)
    output: OutputConfig = Field(default_factory=OutputConfig)
    dataset: list[dict[str, Any]] = Field(default_factory=list)

    @field_validator("dataset")
    @classmethod
    def validate_dataset(cls, value: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
        return [dict(item) for item in value]

PipelineSectionConfig

Bases: SerializableModel

Registered methods and runner behavior.

Source code in src/SafeLens/core/base.py

class PipelineSectionConfig(SerializableModel):
    """Registered methods and runner behavior."""

    probes: list[MethodSpec] = Field(default_factory=list)
    monitors: list[MethodSpec] = Field(default_factory=list)
    attributors: list[MethodSpec] = Field(default_factory=list)
    risk_threshold: float = Field(default=0.5, ge=0.0, le=1.0)

ProbeResult

Bases: SerializableModel

Result returned by an endogenous safety probe.

Source code in src/SafeLens/core/base.py

class ProbeResult(SerializableModel):
    """Result returned by an endogenous safety probe."""

    risk_score: float = Field(ge=0.0, le=1.0)
    critical_layers: list[LayerRef] = Field(default_factory=list)
    intervention_applied: bool = False
    details: dict[str, Any] = Field(default_factory=dict)

RunReport

Bases: SerializableModel

Aggregate report written by a pipeline run.

Source code in src/SafeLens/core/base.py

class RunReport(SerializableModel):
    """Aggregate report written by a pipeline run."""

    generated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    reports: list[SafetyReport] = Field(default_factory=list)
    summary: dict[str, Any] = Field(default_factory=dict)

SafetyReport

Bases: SerializableModel

Standard report shape consumed by downstream adapters such as FlagSafe.

Source code in src/SafeLens/core/base.py

class SafetyReport(SerializableModel):
    """Standard report shape consumed by downstream adapters such as FlagSafe."""

    sample_id: str | None = None
    flagged: bool = False
    risk_score: float = Field(default=0.0, ge=0.0, le=1.0)
    risk_category: list[str] = Field(default_factory=list)
    evidence_tokens: list[int] = Field(default_factory=list)
    attribution_score: float | None = Field(default=None, ge=0.0, le=1.0)
    probe_results: list[ProbeResult] = Field(default_factory=list)
    monitoring_signals: list[MonitoringSignal] = Field(default_factory=list)
    attributions: list[AttributionResult] = Field(default_factory=list)
    metadata: dict[str, Any] = Field(default_factory=dict)

SerializableModel

Bases: BaseModel

Pydantic base model with a stable dictionary export helper.

Source code in src/SafeLens/core/base.py

class SerializableModel(BaseModel):
    """Pydantic base model with a stable dictionary export helper."""

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")

    def to_dict(self) -> dict[str, Any]:
        """Return a JSON-serializable representation."""
        return self.model_dump(mode="json")

to_dict()

Return a JSON-serializable representation.

Source code in src/SafeLens/core/base.py

def to_dict(self) -> dict[str, Any]:
    """Return a JSON-serializable representation."""
    return self.model_dump(mode="json")

TokenAttribution

Bases: SerializableModel

Token-level attribution evidence.

Source code in src/SafeLens/core/base.py

class TokenAttribution(SerializableModel):
    """Token-level attribution evidence."""

    token_index: int = Field(ge=0)
    score: float
    token_text: str | None = None
    source: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)