evalrx 0.1.2__py3-none-any.whl

evalrx/core/experiment.py

@@ -0,0 +1,106 @@
+"""Experiment — a declarative spec that bridges the agent and the core.
+
+The self-evolving agent does not call analyzers directly; it emits *experiments*
+("run analyzer X on model Y over cases Z"), which an :class:`ExperimentRunner`
+executes, caches, and turns into results.  Making the experiment a first-class
+object lets the agent queue, log, replay, and reason about its own actions.
+
+This is intentionally minimal in Stage 1: enough structure for the agent layer
+to target, with caching as the one piece of real behaviour.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from evalrx.core.analyzer import Analyzer
+    from evalrx.core.model import Model
+    from evalrx.core.result import Result
+
+
+def _data_fingerprint(data: Any) -> Any:
+    """Deterministic content fingerprint for cache keys.
+
+    Crucially does NOT call ``as_casebatch`` (which mints fresh uuids for raw
+    inputs) — that would make the same data hash differently each call.  Uses
+    stable case ids when present, otherwise the prompt/content text.
+    """
+    from evalrx.core.case import CaseBatch, FailureCase, Inputs
+
+    if isinstance(data, CaseBatch):
+        return [c.id for c in data]
+    if isinstance(data, FailureCase):
+        return [data.id]
+    if isinstance(data, Inputs):
+        return [f"inputs:{data.prompt}"]
+    if isinstance(data, str):
+        return [f"str:{data}"]
+    if isinstance(data, (list, tuple)):
+        return [_data_fingerprint(x) for x in data]
+    return [repr(data)]
+
+
+@dataclass
+class Experiment:
+    """A declarative analysis request.
+
+    Attributes:
+        model:    The model to analyse.
+        analyzer: The configured analyzer to run.
+        data:     Cases (or anything :func:`as_casebatch` accepts).
+        metadata: Free-form context (e.g. the hypothesis this tests).
+    """
+
+    model: "Model"
+    analyzer: "Analyzer"
+    data: Any
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def fingerprint(self) -> str:
+        """A stable CONTENT hash of (model, analyzer+params, data).
+
+        Used as the cache key so that two *equivalent* experiments dedupe (and a
+        rerun after restart hits cache) — unlike ``id()``, which is a transient
+        memory address.  The data fingerprint uses case ids when available.
+        """
+        try:
+            params = self.analyzer.get_params()
+        except Exception:  # non-estimator analyzer
+            params = {}
+        data_fp = _data_fingerprint(self.data)
+        payload = json.dumps(
+            {
+                "model": repr(self.model),
+                "analyzer": getattr(self.analyzer, "name", type(self.analyzer).__name__),
+                "params": params,
+                "data": data_fp,
+            },
+            sort_keys=True,
+            default=repr,
+        )
+        return hashlib.sha1(payload.encode()).hexdigest()[:16]
+
+
+class ExperimentRunner:
+    """Execute :class:`Experiment` specs, caching results by experiment CONTENT."""
+
+    def __init__(self) -> None:
+        self._cache: dict[str, "Result"] = {}
+
+    def run(self, experiment: Experiment, use_cache: bool = True) -> "Result":
+        """Run an experiment, returning (and caching) its :class:`Result`."""
+        key = experiment.fingerprint()
+        if use_cache and key in self._cache:
+            return self._cache[key]
+        result = experiment.analyzer.run(experiment.model, experiment.data)
+        result.metadata.setdefault("experiment", experiment.metadata)
+        result.metadata.setdefault("fingerprint", key)
+        self._cache[key] = result
+        return result
+
+    def run_many(self, experiments: list[Experiment]) -> list["Result"]:
+        return [self.run(exp) for exp in experiments]

evalrx/core/model.py

@@ -0,0 +1,198 @@
+"""Model — the abstract contract every analyzable model implements.
+
+The key design move: a model exposes **one** rich forward pass,
+``forward(inputs, capture={...}) -> Trace``, where ``capture`` names the
+internals to record (attention, hidden states, …).  Analyzers consume the
+``Trace``, not the model's framework-specific guts.  This decouples "run the
+model and grab internals" from "interpret the internals", so one model serves
+every compatible analyzer and one analyzer works across every compatible model.
+
+The hybrid convenience API (``model.call_attention(prompt)``) is provided here
+by ``__getattr__``: it resolves ``attention`` in the analyzer registry, checks
+the model supports it, and delegates to ``AttentionAnalyzer().run(self, ...)``.
+No mixins, no per-model wiring — the shim is derived from capabilities.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable
+
+from evalrx.core.capability import Capability
+
+if TYPE_CHECKING:
+    import torch
+
+    from evalrx.core.tokentype import TokenTypeMap
+    from evalrx.core.tool import ChatTurn
+
+
+@dataclass
+class CaptureSpec:
+    """Optional knobs that bound what ``forward`` records, to control memory.
+
+    Capturing every layer/head of attention for a long sequence is O(L^2) and
+    can be tens of GB.  A backend that honours this spec captures only the
+    requested layers/heads and can keep tensors off the CPU.  ``None`` fields
+    mean "everything" (the simple default).  Backends that don't support
+    sub-selection ignore it — the ``capture`` set alone still works.
+    """
+
+    layers: "list[int] | None" = None    # which decoder layers to keep (None = all)
+    heads: "list[int] | None" = None     # which attention heads to keep (None = all)
+    to_cpu: bool = True                   # move captured tensors to CPU (frees GPU mem)
+    with_grad: bool = False               # keep the graph for GRADIENTS (else no_grad)
+
+
+@dataclass
+class TokenLogprob:
+    """Logprob of one output token (+ optional top-k alternatives).
+
+    The black-box uncertainty primitive: retrievable from an API that returns
+    ``logprobs`` (OpenAI) — no model internals needed.
+    """
+
+    token: str
+    logprob: float
+    top: dict[str, float] = field(default_factory=dict)  # token -> logprob (top-k alternatives)
+
+
+@dataclass
+class Trace:
+    """Captured internals from a single forward pass.
+
+    Only the fields named in the ``capture`` set are populated; the rest stay
+    ``None``.  ``provided`` records what was actually captured.
+    """
+
+    tokens: list[str]
+    token_ids: list[int]
+    provided: set[Capability] = field(default_factory=set)
+    attentions: "list[torch.Tensor] | None" = None      # per layer: (heads, seq, seq)
+    hidden_states: "list[torch.Tensor] | None" = None    # per layer: (seq, dim)
+    logits: "torch.Tensor | None" = None                 # (seq, vocab)
+    token_type_map: "TokenTypeMap | None" = None         # VLM: image-token positions + patch grid
+    extras: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def seq_len(self) -> int:
+        return len(self.tokens)
+
+    def require(self, capability: Capability) -> Any:
+        """Return the field for *capability*, raising if it wasn't captured."""
+        attr = {
+            Capability.ATTENTION: "attentions",
+            Capability.HIDDEN_STATES: "hidden_states",
+            Capability.LOGITS: "logits",
+        }.get(capability)
+        value = getattr(self, attr) if attr else None
+        if value is None:
+            raise ValueError(
+                f"Trace does not contain {capability.value!r}. "
+                f"Was it included in the forward(capture=...) set? "
+                f"Captured: {sorted(c.value for c in self.provided)}"
+            )
+        return value
+
+
+class Model(ABC):
+    """Abstract base for every analyzable model.
+
+    Subclasses must:
+      1. set the ``capabilities`` class attribute,
+      2. implement :meth:`generate` and :meth:`forward`.
+
+    They get ``call_<analyzer>(data, **kwargs)`` for free via :meth:`__getattr__`.
+    """
+
+    #: Capabilities this model provides. Override in subclasses.
+    capabilities: frozenset[Capability] = frozenset()
+
+    #: Modalities this model handles (e.g. {"text"}, {"text","image"}). Analyzers
+    #: declare ``applies_to_modalities`` and are matched against this.
+    modalities: frozenset[str] = frozenset({"text"})
+
+    def unembed_weight(self):
+        """Optional: the unembedding matrix ``(vocab, dim)`` for logit-lens analyzers.
+
+        Returns ``None`` unless a white-box backend exposes it.
+        """
+        return None
+
+    @abstractmethod
+    def generate(self, inputs: Any, **kwargs) -> str:
+        """Produce a text response for *inputs* (a prompt or :class:`Inputs`)."""
+
+    @abstractmethod
+    def forward(
+        self,
+        inputs: Any,
+        capture: set[Capability],
+        spec: "CaptureSpec | None" = None,
+    ) -> Trace:
+        """Run one forward pass, capturing the requested internals into a :class:`Trace`.
+
+        ``capture`` names *what* to record; the optional ``spec`` bounds *how
+        much* (layers/heads/device) for memory control.  Backends free to ignore
+        ``spec`` when they cannot sub-select.
+        """
+
+
+    def chat(self, messages: list, tools: "list | None" = None) -> "ChatTurn":
+        """One tool-aware turn for agent mode (the ``Agent`` loop calls this).
+
+        Returns a :class:`~evalrx.core.tool.ChatTurn` (assistant text + any
+        native tool_calls).  Only tool-capable handles override this; the default
+        signals that this model cannot drive an agent loop.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement chat(); it is not TOOL_CALLS-capable."
+        )
+
+    def logprobs(self, inputs: Any, **kwargs) -> "list[TokenLogprob]":
+        """Generate and return per-output-token logprobs (the LOGPROBS capability).
+
+        Black-box-feasible (OpenAI-style ``logprobs``).  Only LOGPROBS-capable
+        handles override this.
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not implement logprobs(); it is not LOGPROBS-capable."
+        )
+
+
+    # ------------------------------------------------------------------
+    # Capability checks
+    # ------------------------------------------------------------------
+
+    def supports(self, required: frozenset[Capability] | set[Capability]) -> bool:
+        """True iff this model provides every capability in *required*."""
+        return set(required).issubset(self.capabilities)
+
+    # ------------------------------------------------------------------
+    # Hybrid convenience shim:  model.call_<analyzer>(data, **kwargs)
+    # ------------------------------------------------------------------
+
+    def __getattr__(self, attr: str) -> Callable[..., Any]:
+        # __getattr__ only fires for attributes not found normally, so this
+        # never shadows real methods/fields.
+        if attr.startswith("call_"):
+            analyzer_name = attr[len("call_"):]
+            from evalrx.core.registry import registry
+
+            if registry.analyzers.has(analyzer_name):
+                analyzer_cls = registry.analyzers.get(analyzer_name)
+
+                def _shim(data: Any, **kwargs) -> Any:
+                    return analyzer_cls(**kwargs).run(self, data)
+
+                _shim.__name__ = attr
+                _shim.__doc__ = (
+                    f"Convenience shim → {analyzer_cls.__name__}(**kwargs).run(self, data). "
+                    f"Auto-derived from capabilities + registry."
+                )
+                return _shim
+
+        raise AttributeError(
+            f"{type(self).__name__!r} object has no attribute {attr!r}"
+        )

evalrx/core/pipeline.py

@@ -0,0 +1,42 @@
+"""Pipeline — sklearn-style composition of analyzers.
+
+Run several analyzers over the same ``(model, cases)`` and collect their
+results under their step names::
+
+    pipe = Pipeline([
+        ("attention", AttentionAnalyzer(layer=-1)),
+        # ("saliency", SaliencyAnalyzer()),   # Stage 2
+    ])
+    results = pipe.run(qwen, "The capital of France is")
+    results["attention"].summary()
+
+Each analyzer's capabilities are validated independently by its own ``run``.
+Cross-step data flow (e.g. perturb → analyze) is a Stage-2 extension.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from evalrx.core.case import as_casebatch
+
+if TYPE_CHECKING:
+    from evalrx.core.analyzer import Analyzer
+    from evalrx.core.model import Model
+    from evalrx.core.result import Result
+
+
+class Pipeline:
+    """Ordered collection of named analyzers run over shared inputs."""
+
+    def __init__(self, steps: list[tuple[str, "Analyzer"]]) -> None:
+        self.steps = steps
+
+    def run(self, model: "Model", data: Any) -> dict[str, "Result"]:
+        """Run every step over the same normalised case batch."""
+        cases = as_casebatch(data)
+        return {name: analyzer.run(model, cases) for name, analyzer in self.steps}
+
+    def __repr__(self) -> str:
+        names = ", ".join(name for name, _ in self.steps)
+        return f"Pipeline([{names}])"

evalrx/core/registry.py

@@ -0,0 +1,142 @@
+"""Registries — the agent's discovery surface.
+
+Analyzers register themselves by name.  An agent (or the framework) enumerates
+them and matches them to a model by capability::
+
+    registry.analyzers.list()                  # ["attention", "token_entropy", ...]
+    registry.analyzers.compatible_with(model)  # [AttentionAnalyzer, ...]
+
+Model *identity* lives in :mod:`evalrx.specs` (``ModelSpec``), not here;
+``registry.models`` is a deprecated shim that delegates to it.  This is what
+makes the package programmatically explorable rather than something you have to
+read source code to use.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from evalrx.core.analyzer import Analyzer
+    from evalrx.core.model import Model
+
+T = TypeVar("T")
+
+
+class _Registry(Generic[T]):
+    """Name → class mapping with a registration decorator."""
+
+    def __init__(self, kind: str) -> None:
+        self._kind = kind
+        self._items: dict[str, type[T]] = {}
+
+    def register(self, name: str) -> Callable[[type[T]], type[T]]:
+        """Decorator: ``@registry.register("qwen")``."""
+
+        def decorator(cls: type[T]) -> type[T]:
+            key = name.lower()
+            if key in self._items:
+                raise ValueError(f"{self._kind} '{name}' is already registered.")
+            self._items[key] = cls
+            return cls
+
+        return decorator
+
+    def get(self, name: str) -> type[T]:
+        key = name.lower()
+        if key not in self._items:
+            raise KeyError(
+                f"Unknown {self._kind} '{name}'. Registered: {self.list()}"
+            )
+        return self._items[key]
+
+    def has(self, name: str) -> bool:
+        return name.lower() in self._items
+
+    def list(self) -> list[str]:
+        return sorted(self._items)
+
+    def all(self) -> dict[str, type[T]]:
+        return dict(self._items)
+
+
+class AnalyzerRegistry(_Registry["Analyzer"]):
+    """Analyzer registry with capability-based matching."""
+
+    @staticmethod
+    def _matches(cls: type["Analyzer"], model: "Model") -> bool:
+        """Analyzer runs on *model* iff capabilities are satisfied AND modalities overlap."""
+        if not model.supports(cls.requires):
+            return False
+        model_modalities = getattr(model, "modalities", frozenset({"text"}))
+        return bool(cls.applies_to_modalities & model_modalities)
+
+    def compatible_with(self, model: "Model") -> list[type["Analyzer"]]:
+        """Return analyzer classes runnable on *model* (capability + modality match)."""
+        return [cls for cls in self._items.values() if self._matches(cls, model)]
+
+    def names_compatible_with(self, model: "Model") -> list[str]:
+        """Same as :meth:`compatible_with` but returns registered names."""
+        return sorted(name for name, cls in self._items.items() if self._matches(cls, model))
+
+
+class _DeprecatedModelRegistry(_Registry["Model"]):
+    """Deprecated shim: model identity now lives in :mod:`evalrx.specs`.
+
+    ``registry.models`` is kept so older code keeps working, but it delegates to
+    the spec registry and warns.  Use ``evalrx.list_specs()`` /
+    ``evalrx.get_spec()`` (or ``evalrx.load`` / ``compose``) instead.
+    """
+
+    @staticmethod
+    def _warn() -> None:
+        import warnings
+
+        warnings.warn(
+            "registry.models is deprecated; model identity now lives in "
+            "evalrx.specs (use evalrx.list_specs() / get_spec(), or "
+            "evalrx.load / compose to build a model).",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+
+    def list(self) -> list[str]:
+        self._warn()
+        from evalrx.specs import list_specs
+
+        return list_specs()
+
+    def has(self, name: str) -> bool:
+        self._warn()
+        from evalrx.specs import REGISTRY
+
+        return name.lower() in REGISTRY
+
+    def get(self, name: str):  # returns a ModelSpec, not a class
+        self._warn()
+        from evalrx.specs import get_spec
+
+        return get_spec(name.lower())
+
+
+class Registry:
+    """Top-level namespace bundling the model and analyzer registries."""
+
+    def __init__(self) -> None:
+        # Model identity is the spec registry's job now; this is a back-compat shim.
+        self.models: _DeprecatedModelRegistry = _DeprecatedModelRegistry("model")
+        self.analyzers: AnalyzerRegistry = AnalyzerRegistry("analyzer")
+
+
+# Singleton used across the package.
+registry = Registry()
+
+
+def register_model(name: str):
+    """Class decorator registering a model under *name*."""
+    return registry.models.register(name)
+
+
+def register_analyzer(name: str):
+    """Class decorator registering an analyzer under *name*."""
+    return registry.analyzers.register(name)

evalrx/core/result.py

@@ -0,0 +1,64 @@
+"""Result — the structured, agent-readable output of every analyzer.
+
+A result has two faces:
+  - ``findings``: a light, JSON-serialisable dict that an LLM agent reads to
+    decide what to do next (e.g. top attended tokens, entropy stats).
+  - ``artifacts``: heavy objects (attention tensors, hidden states) for humans,
+    plotting, and downstream numeric analysis — never sent to the agent verbatim.
+
+Subclasses (e.g. ``AttentionResult``) add domain-specific convenience methods
+backed by ``artifacts`` while keeping the uniform ``findings`` contract.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from evalrx.core.case import CaseBatch
+
+
+@dataclass
+class Result:
+    """Uniform analyzer output.
+
+    Attributes:
+        analyzer:  Registered name of the analyzer that produced this.
+        model:     ``repr()`` of the analysed model.
+        findings:  Light, JSON-serialisable summary (the agent-facing answer).
+        artifacts: Heavy objects (tensors, arrays) — not serialised by default.
+        cases:     The :class:`~evalrx.core.case.CaseBatch` analysed.
+        metadata:  Free-form extra info.
+    """
+
+    analyzer: str
+    model: str
+    findings: dict[str, Any] = field(default_factory=dict)
+    artifacts: dict[str, Any] = field(default_factory=dict)
+    cases: "CaseBatch | None" = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def summary(self) -> str:
+        """Human/LLM-readable one-screen summary built from ``findings``."""
+        lines = [f"[{self.analyzer}] on {self.model}"]
+        for key, value in self.findings.items():
+            lines.append(f"  {key}: {value}")
+        return "\n".join(lines)
+
+    def to_dict(self) -> dict[str, Any]:
+        """JSON-serialisable view (``findings`` only — artifacts are dropped)."""
+        return {
+            "analyzer": self.analyzer,
+            "model": self.model,
+            "findings": self.findings,
+            "metadata": self.metadata,
+            "n_cases": len(self.cases) if self.cases is not None else 0,
+        }
+
+    def to_json(self, **kwargs) -> str:
+        return json.dumps(self.to_dict(), **kwargs)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(analyzer={self.analyzer!r}, model={self.model!r})"