pfnstudio-core 0.7.0__tar.gz

pfnstudio_core-0.7.0/.gitignore

@@ -0,0 +1,68 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+.venv/
+.deps-venv/
+venv/
+ENV/
+env/
+
+# Editable-install metadata
+_editable_impl_*.pth
+
+# IDE / editor
+.vscode/
+.idea/
+*.swp
+*.swo
+.DS_Store
+
+# pytest / coverage / mypy
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+.coverage.*
+htmlcov/
+.tox/
+.nox/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Build artifacts
+checkpoint/
+runs/*/checkpoint/
+runs/*/results.json
+runs/*/metrics.jsonl
+*.pt
+*.pth
+
+# Dataset cache (downloaded benchmark datasets)
+.priorstudio/datasets/
+
+# Logs
+logs/
+*.log

pfnstudio_core-0.7.0/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: pfnstudio-core
+Version: 0.7.0
+Summary: Core abstractions for PFN Studio — Prior, Model, Eval, Run, and the block registry.
+Project-URL: Homepage, https://pfnstudio.com
+Project-URL: Repository, https://github.com/profitopsai/pfnstudio
+Project-URL: Issues, https://github.com/profitopsai/pfnstudio/issues
+Project-URL: Documentation, https://github.com/profitopsai/pfnstudio/tree/main/docs
+Author: PFN Studio contributors
+License: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.26
+Requires-Dist: pandas>=2.0
+Requires-Dist: pydantic>=2.6
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: scikit-learn>=1.3
+Provides-Extra: dev
+Requires-Dist: pytest-cov>=5.0; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Provides-Extra: torch
+Requires-Dist: torch>=2.2; extra == 'torch'
+Description-Content-Type: text/markdown
+
+# pfnstudio-core
+
+The Python contract for PFN Studio FM projects.
+
+```python
+from pfnstudio_core import Prior, Model, Eval, Run, register_block, register_prior
+
+@register_prior("my_prior")
+class MyPrior(Prior):
+    def sample(self, seed: int): ...
+
+@register_block("my_attention")
+class MyAttention:
+    def __init__(self, d_model: int, n_heads: int): ...
+```
+
+The CLI discovers anything registered via these decorators and validates `models/*.yaml` references against the registry.
+
+## Layout
+
+- `prior.py` — `Prior` ABC and built-in prior loader
+- `model.py` — `Model` config + block-composition
+- `eval.py` — `Eval` config + result schema
+- `run.py` — `Run` manifest + executor protocol
+- `registry.py` — `@register_prior`, `@register_block`, `@register_eval` and discovery
+- `loaders.py` — load YAML artifacts into typed objects
+- `blocks/` — built-in architecture blocks (transformer encoder, causal attention, heads)
+- `training/` — minimal in-process training loop for the `local` compute adapter

pfnstudio_core-0.7.0/README.md

@@ -0,0 +1,28 @@
+# pfnstudio-core
+
+The Python contract for PFN Studio FM projects.
+
+```python
+from pfnstudio_core import Prior, Model, Eval, Run, register_block, register_prior
+
+@register_prior("my_prior")
+class MyPrior(Prior):
+    def sample(self, seed: int): ...
+
+@register_block("my_attention")
+class MyAttention:
+    def __init__(self, d_model: int, n_heads: int): ...
+```
+
+The CLI discovers anything registered via these decorators and validates `models/*.yaml` references against the registry.
+
+## Layout
+
+- `prior.py` — `Prior` ABC and built-in prior loader
+- `model.py` — `Model` config + block-composition
+- `eval.py` — `Eval` config + result schema
+- `run.py` — `Run` manifest + executor protocol
+- `registry.py` — `@register_prior`, `@register_block`, `@register_eval` and discovery
+- `loaders.py` — load YAML artifacts into typed objects
+- `blocks/` — built-in architecture blocks (transformer encoder, causal attention, heads)
+- `training/` — minimal in-process training loop for the `local` compute adapter

pfnstudio_core-0.7.0/pfnstudio_core/__init__.py

@@ -0,0 +1,68 @@
+"""PFN Studio core abstractions."""
+
+from . import axes as _axes  # registers built-in axes on import
+from . import blocks as _blocks  # registers built-in blocks on import
+from .axes import (
+    UNKNOWN,
+    Axis,
+    AxisDetector,
+    detect,
+    encode_tag,
+    get_axis,
+    is_unknown,
+    list_axes,
+    register_axis,
+    sample_tag,
+    tag_dim,
+    train_detector,
+)
+from .datasets import DatasetUnavailable, RegistryDatasetLoader
+from .eval import Eval, EvalSpec
+from .model import Model, ModelSpec
+from .prior import Prior, PriorSpec
+from .registry import (
+    get_block,
+    get_eval,
+    get_prior,
+    list_blocks,
+    list_evals,
+    list_priors,
+    register_block,
+    register_eval,
+    register_prior,
+)
+from .run import Run, RunSpec
+
+__version__ = "0.4.0"
+
+__all__ = [
+    "UNKNOWN",
+    "Axis",
+    "AxisDetector",
+    "Eval",
+    "EvalSpec",
+    "Model",
+    "ModelSpec",
+    "Prior",
+    "PriorSpec",
+    "Run",
+    "RunSpec",
+    "detect",
+    "encode_tag",
+    "get_axis",
+    "get_block",
+    "get_eval",
+    "get_prior",
+    "is_unknown",
+    "list_axes",
+    "list_blocks",
+    "list_evals",
+    "list_priors",
+    "register_axis",
+    "register_block",
+    "register_eval",
+    "register_prior",
+    "sample_tag",
+    "tag_dim",
+    "train_detector",
+]

pfnstudio_core-0.7.0/pfnstudio_core/axes/__init__.py

@@ -0,0 +1,34 @@
+"""Axes package — public surface for promptable-prior axes.
+
+Importing this module also imports `builtins`, which eagerly
+registers the built-in axes (monotonicity, …) so they're available
+without an extra import in user code."""
+
+from . import builtins as _builtins  # noqa: F401  (side-effect: registers built-ins)
+from .base import (
+    UNKNOWN,
+    Axis,
+    AxisKind,
+    get_axis,
+    is_unknown,
+    list_axes,
+    register_axis,
+)
+from .detector import AxisDetector, detect, train_detector
+from .encoding import encode_tag, sample_tag, tag_dim
+
+__all__ = [
+    "UNKNOWN",
+    "Axis",
+    "AxisDetector",
+    "AxisKind",
+    "detect",
+    "encode_tag",
+    "get_axis",
+    "is_unknown",
+    "list_axes",
+    "register_axis",
+    "sample_tag",
+    "tag_dim",
+    "train_detector",
+]

pfnstudio_core-0.7.0/pfnstudio_core/axes/base.py

@@ -0,0 +1,111 @@
+"""Axis abstraction for promptable priors.
+
+An Axis declares a *steerable property* of a prior — a knob the trained
+brain learns to honor at inference time. Examples: monotonicity (edge
+signs in an SCM), lag_scale (delay distribution), feedback_allowed
+(DAG vs cyclic), sparsity (mean parents per node).
+
+The axis declares the contract (name, kind, value space, default).
+Each Prior that opts into an axis implements its own application
+logic — the same axis can mean different things to different priors
+(e.g. monotonicity on an SCM constrains edge signs; on a regression
+prior it constrains weight signs). The base class doesn't try to
+unify those interpretations.
+
+Back-compat invariants this module enforces:
+- The reserved sentinel ``UNKNOWN`` means "skip this axis at sample
+  time". A prior given ``tag={axis: UNKNOWN}`` for every axis is
+  bit-identical to its pre-axis behavior. This is what guarantees
+  adding axes can't regress existing benchmarks.
+- A prior with no declared axes ignores any ``tag=...`` argument
+  passed to ``sample()`` — old call sites keep working unchanged.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+AxisKind = Literal["categorical", "range", "boolean"]
+
+# Reserved sentinel — when a tag has this for an axis, the prior must
+# skip the axis hook and behave as if the axis weren't declared.
+# Stored as a string so it round-trips through JSON / YAML cleanly.
+UNKNOWN = "__unknown__"
+
+
+@dataclass(frozen=True)
+class Axis:
+    name: str
+    kind: AxisKind
+    # For categorical: list of label strings. For range: [min, max].
+    # For boolean: ignored.
+    values: tuple[Any, ...] = field(default_factory=tuple)
+    description: str = ""
+    # Fraction of training samples that get UNKNOWN for this axis.
+    # Default 0.3 — preserves substantial unconditional mass so the
+    # brain doesn't lose its pre-axis baseline.
+    unknown_mass: float = 0.3
+
+    def __post_init__(self) -> None:
+        if not 0.0 <= self.unknown_mass <= 1.0:
+            raise ValueError(
+                f"axis {self.name!r}: unknown_mass must be in [0, 1], got {self.unknown_mass}"
+            )
+        if self.kind == "categorical" and not self.values:
+            raise ValueError(f"axis {self.name!r}: categorical axes need at least one value")
+        if self.kind == "range" and len(self.values) != 2:
+            raise ValueError(
+                f"axis {self.name!r}: range axes need exactly 2 values (min, max), got {len(self.values)}"
+            )
+
+    def sample_value(self, rng: Any) -> Any:
+        """Sample one value for this axis, honoring the unknown_mass invariant."""
+        if rng.random() < self.unknown_mass:
+            return UNKNOWN
+        if self.kind == "categorical":
+            return rng.choice(list(self.values))
+        if self.kind == "boolean":
+            return bool(rng.integers(0, 2))
+        if self.kind == "range":
+            lo, hi = float(self.values[0]), float(self.values[1])
+            return float(rng.uniform(lo, hi))
+        raise ValueError(f"unknown axis kind: {self.kind!r}")
+
+
+# Module-level registry. Resolves axis names → Axis instances at
+# loader time. Built-in axes register themselves on import; custom
+# axes register via @register_axis.
+_AXES: dict[str, Axis] = {}
+
+
+def register_axis(axis: Axis) -> Axis:
+    """Register an Axis. Idempotent: re-registering the *same* axis is
+    a no-op (useful when modules get imported twice); re-registering
+    a *different* axis under the same name raises."""
+    existing = _AXES.get(axis.name)
+    if existing is not None and existing != axis:
+        raise ValueError(
+            f"axis {axis.name!r} already registered with a different definition"
+        )
+    _AXES[axis.name] = axis
+    return axis
+
+
+def get_axis(name: str) -> Axis:
+    """Look up a registered axis by name. Raises KeyError if not registered."""
+    if name not in _AXES:
+        raise KeyError(
+            f"axis {name!r} not registered. Known axes: {sorted(_AXES)}"
+        )
+    return _AXES[name]
+
+
+def list_axes() -> list[str]:
+    return sorted(_AXES)
+
+
+def is_unknown(value: Any) -> bool:
+    """Whether a tag value is the UNKNOWN sentinel. Use this in axis
+    application code rather than comparing to the string literal."""
+    return value == UNKNOWN

pfnstudio_core-0.7.0/pfnstudio_core/axes/builtins/__init__.py

@@ -0,0 +1,5 @@
+"""Built-in axes. Importing this module registers them with the global
+axis registry. The package ``__init__`` imports this so axes are
+available the moment ``pfnstudio_core`` is imported."""
+
+from . import monotonicity  # noqa: F401  (registers on import)

pfnstudio_core-0.7.0/pfnstudio_core/axes/builtins/monotonicity.py

@@ -0,0 +1,25 @@
+"""Built-in `monotonicity` axis.
+
+Declares the contract; per-prior application lives in each Prior
+subclass (an SCM's monotonicity constrains edge signs, a regression
+prior's monotonicity constrains weight signs, etc.).
+"""
+
+from __future__ import annotations
+
+from ..base import Axis, register_axis
+
+monotonicity = register_axis(
+    Axis(
+        name="monotonicity",
+        kind="categorical",
+        values=("positive", "negative", "mixed"),
+        description=(
+            "Sign discipline of the prior's causal relationships. "
+            "positive = X up implies Y up everywhere; negative = "
+            "X up implies Y down; mixed = signs may vary. Unknown "
+            "(default) leaves the prior unconstrained."
+        ),
+        unknown_mass=0.3,
+    )
+)

pfnstudio_core-0.7.0/pfnstudio_core/axes/detector.py

@@ -0,0 +1,260 @@
+"""Auto-detector — supervised model that reads context and proposes
+axis values.
+
+The pitch in one line: instead of asking the user to fill out an empty
+form of axis chips, the detector reads their data and pre-fills the
+chips with its best guess (and a confidence). The user just confirms,
+overrides, or adds the axes the detector can't see.
+
+Training:
+- Each batch, the prior samples a real (non-UNKNOWN) tag value and
+  generates data from it. The tag value is the supervised label.
+- The detector reads the data (no tag) and predicts a softmax over
+  axis values.
+- Cross-entropy loss.
+
+Inference:
+- Same encoder forward, output softmax probabilities.
+- Highest-probability value becomes the proposed chip; that probability
+  becomes the displayed confidence.
+
+Limitations in v1:
+- Only categorical axes are detected. Range axes need a regression
+  head and a different loss; boolean axes need a 2-class head.
+  Both are straightforward additions but defer until they're actually
+  shipped on a prior.
+- The detector is small (32-d, 2-layer) — same shape as the test
+  brain. Scale up when a real prior with many axes ships.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .base import UNKNOWN, Axis
+
+
+class AxisDetector:
+    """A small encoder + per-axis classification head. Holds one
+    head per categorical axis; non-categorical axes are silently
+    skipped (an empty detector is legal — predicts nothing)."""
+
+    def __init__(
+        self,
+        axes: list[Axis],
+        d_model: int = 32,
+        n_heads: int = 4,
+        n_layers: int = 2,
+        dropout: float = 0.0,
+    ):
+        try:
+            import torch.nn as nn
+        except ImportError as e:
+            raise ImportError(
+                "AxisDetector requires torch. "
+                "Install with: pip install pfnstudio-core[torch]"
+            ) from e
+
+        # Only categorical axes get heads in v1. Skipped axes are
+        # tracked so callers can warn the user "we can't propose for
+        # axis X yet."
+        self._scored_axes: list[Axis] = [
+            a for a in axes if a.kind == "categorical" and len(a.values) >= 2
+        ]
+        self._skipped_axes: list[Axis] = [
+            a for a in axes if a not in self._scored_axes
+        ]
+        self.d_model = d_model
+
+        # Lazy embedder — input feature dim is inferred on first
+        # forward. Same pattern the TabularEmbedder uses.
+        self.embedder = nn.LazyLinear(d_model)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=n_heads,
+            dim_feedforward=d_model * 4,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
+        # One head per scored axis. Output dim = number of values
+        # (no UNKNOWN class — the detector proposes a real value;
+        # low confidence across all classes is how it expresses
+        # "I'm not sure", not a separate class).
+        self.heads = nn.ModuleDict(
+            {a.name: nn.Linear(d_model, len(a.values)) for a in self._scored_axes}
+        )
+
+    @property
+    def scored_axes(self) -> list[Axis]:
+        """Axes the detector has heads for (categorical, ≥2 values)."""
+        return self._scored_axes
+
+    @property
+    def skipped_axes(self) -> list[Axis]:
+        """Axes the detector can't predict (boolean, range, single-value)."""
+        return self._skipped_axes
+
+    def parameters(self):
+        # nn.ModuleDict + nn.Module attributes — yield everything.
+        yield from self.embedder.parameters()
+        yield from self.encoder.parameters()
+        yield from self.heads.parameters()
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "embedder": self.embedder.state_dict(),
+            "encoder": self.encoder.state_dict(),
+            "heads": self.heads.state_dict(),
+        }
+
+    def __call__(self, x: Any) -> dict[str, Any]:
+        """Forward pass. Returns axis_name → logits tensor of shape
+        (B, num_values_for_axis)."""
+        emb = self.embedder(x)  # (B, N, d_model)
+        encoded = self.encoder(emb)  # (B, N, d_model)
+        # Mean-pool the context tokens to one vector per batch element.
+        # (Could swap for an attention pool later; mean is fine for
+        # the monotonicity demo and avoids extra parameters.)
+        pooled = encoded.mean(dim=1)  # (B, d_model)
+        return {name: head(pooled) for name, head in self.heads.items()}
+
+
+def _value_to_index(axis: Axis, value: Any) -> int:
+    """Index of `value` in axis.values. Raises if value is UNKNOWN
+    or not a known value — UNKNOWN tags shouldn't be used as training
+    labels because the detector is supposed to learn to assign a
+    concrete value."""
+    if value == UNKNOWN:
+        raise ValueError(
+            f"axis {axis.name!r}: cannot train detector on UNKNOWN label"
+        )
+    return list(axis.values).index(value)
+
+
+def train_detector(
+    *,
+    detector: AxisDetector,
+    prior: Any,
+    steps: int = 500,
+    batch_size: int = 16,
+    lr: float = 1e-3,
+    seed: int = 0,
+    on_step: Any = None,
+) -> dict[str, Any]:
+    """Train the detector on labeled batches sampled from the prior.
+
+    Each batch picks a non-UNKNOWN value for each scored axis, samples
+    data with that tag, and trains the detector to predict the tag
+    from the data. Cross-entropy loss, summed across axes.
+
+    Returns ``{status, final_loss, mean_acc}`` where ``mean_acc`` is
+    the accuracy of the detector on the *last* batch (cheap sanity
+    check that training did something)."""
+    try:
+        import torch
+        import torch.nn.functional as F
+    except ImportError:
+        return {"status": "skipped", "reason": "torch not installed"}
+
+    if not detector.scored_axes:
+        return {"status": "skipped", "reason": "detector has no scored axes"}
+
+    rng = np.random.default_rng(seed)
+    optim = torch.optim.AdamW(list(detector.parameters()), lr=lr)
+
+    last_acc: dict[str, float] = {}
+    final_loss = 0.0
+    for step in range(steps):
+        # Sample one tag value per axis (random, non-UNKNOWN — that's
+        # the supervised label the detector learns to recover).
+        tag = {
+            a.name: rng.choice(list(a.values)) for a in detector.scored_axes
+        }
+        batch = prior.sample_batch(
+            batch_size=batch_size, seed=seed + step * batch_size, tag=tag
+        )
+        X = torch.stack([torch.from_numpy(b["X"]).float() for b in batch])
+
+        logits = detector(X)  # dict axis → (B, num_values)
+        loss = torch.zeros(1)
+        for axis in detector.scored_axes:
+            label_idx = _value_to_index(axis, tag[axis.name])
+            labels = torch.full((X.shape[0],), label_idx, dtype=torch.long)
+            loss = loss + F.cross_entropy(logits[axis.name], labels)
+
+        optim.zero_grad()
+        loss.backward()
+        optim.step()
+
+        final_loss = float(loss.item())
+        if on_step is not None:
+            on_step(step, final_loss)
+
+        # Final-batch accuracy (per axis) — cheap and informative.
+        if step == steps - 1:
+            with torch.no_grad():
+                for axis in detector.scored_axes:
+                    preds = logits[axis.name].argmax(dim=-1)
+                    label_idx = _value_to_index(axis, tag[axis.name])
+                    last_acc[axis.name] = float((preds == label_idx).float().mean().item())
+
+    return {
+        "status": "ok",
+        "steps": steps,
+        "final_loss": final_loss,
+        "last_batch_accuracy": last_acc,
+    }
+
+
+def detect(
+    *,
+    detector: AxisDetector,
+    context: Any,
+) -> dict[str, dict[str, Any]]:
+    """Run the detector on a context tensor and return a tag-shaped
+    proposal with per-axis confidence.
+
+    Returns ``{axis_name: {value, confidence, probs}}`` where:
+      - ``value`` is the predicted axis value (argmax)
+      - ``confidence`` is the softmax probability of that value
+      - ``probs`` is the full {value: probability} dict, for callers
+        that want to render uncertainty (e.g. show "60% positive,
+        35% mixed, 5% negative" instead of just "positive 60%").
+    """
+    try:
+        import torch
+        import torch.nn.functional as F
+    except ImportError:
+        return {}
+
+    if not detector.scored_axes:
+        return {}
+
+    if not torch.is_tensor(context):
+        context = torch.from_numpy(np.asarray(context)).float()
+    if context.dim() == 2:
+        context = context.unsqueeze(0)  # add batch dim
+
+    with torch.no_grad():
+        logits = detector(context)
+
+    out: dict[str, dict[str, Any]] = {}
+    for axis in detector.scored_axes:
+        probs = F.softmax(logits[axis.name], dim=-1)
+        # Mean across batch in case the caller passed multiple contexts —
+        # the detector returns one proposal per axis, not per-sample.
+        mean_probs = probs.mean(dim=0)
+        best_idx = int(mean_probs.argmax().item())
+        out[axis.name] = {
+            "value": axis.values[best_idx],
+            "confidence": float(mean_probs[best_idx].item()),
+            "probs": {
+                axis.values[i]: float(mean_probs[i].item())
+                for i in range(len(axis.values))
+            },
+        }
+    return out