json2vec 0.4.7__tar.gz → 0.4.9__tar.gz

{json2vec-0.4.7/src/json2vec.egg-info → json2vec-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.7
+Version: 0.4.9
 Summary: Schema-first PyTorch models for hierarchical / nested / sequence data structures
 License-Expression: Apache-2.0
 Requires-Python: >=3.12
@@ -14,7 +14,6 @@ Requires-Dist: pydantic>=2.11.7
 Requires-Dist: jmespath>=1.0.1
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: anytree>=2.13.0
-Requires-Dist: ordered-set>=4.1.0
 Requires-Dist: pyarrow>=21.0.0
 Requires-Dist: polars>=1.35.2
 Requires-Dist: numpy>=2.2.6
@@ -22,16 +21,20 @@ Requires-Dist: lightning>=2.6.4
 Requires-Dist: tensordict>=0.10.0
 Requires-Dist: torch>=2.7.1
 Provides-Extra: serving
-Requires-Dist: litserve>=0.2.13; extra == "serving"
+Requires-Dist: fastapi>=0.124.0; extra == "serving"
+Requires-Dist: orjson>=3.10.0; extra == "serving"
 Requires-Dist: pydantic-settings>=2.10.1; extra == "serving"
+Requires-Dist: uvicorn>=0.38.0; extra == "serving"
 Provides-Extra: text
 Requires-Dist: transformers>=4.55.0; extra == "text"
 Provides-Extra: docs
-Requires-Dist: litserve>=0.2.13; extra == "docs"
+Requires-Dist: fastapi>=0.124.0; extra == "docs"
 Requires-Dist: mkdocs-material>=9.6; extra == "docs"
 Requires-Dist: mkdocs-jupyter>=0.26.3; extra == "docs"
 Requires-Dist: mkdocstrings[python]>=0.27; extra == "docs"
+Requires-Dist: orjson>=3.10.0; extra == "docs"
 Requires-Dist: pydantic-settings>=2.10.1; extra == "docs"
+Requires-Dist: uvicorn>=0.38.0; extra == "docs"
 Dynamic: license-file
 
 <h1 align="center"><code>json2vec</code></h1>
@@ -314,7 +317,7 @@ uv sync --extra docs
 ```
 
 The `text` extra installs Hugging Face `transformers`. The `serving` extra
-installs LitServe-backed deployment dependencies. The `docs` extra installs the
+installs FastAPI-backed deployment dependencies. The `docs` extra installs the
 MkDocs toolchain.
 
 ## Documentation Map

{json2vec-0.4.7 → json2vec-0.4.9}/README.md

@@ -278,7 +278,7 @@ uv sync --extra docs
 ```
 
 The `text` extra installs Hugging Face `transformers`. The `serving` extra
-installs LitServe-backed deployment dependencies. The `docs` extra installs the
+installs FastAPI-backed deployment dependencies. The `docs` extra installs the
 MkDocs toolchain.
 
 ## Documentation Map

{json2vec-0.4.7 → json2vec-0.4.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "json2vec"
-version = "0.4.7"
+version = "0.4.9"
 description = "Schema-first PyTorch models for hierarchical / nested / sequence data structures"
 readme = "README.md"
 license = "Apache-2.0"
@@ -13,7 +13,6 @@ dependencies = [
     "jmespath>=1.0.1",
     "loguru>=0.7.3",
     "anytree>=2.13.0",
-    "ordered-set>=4.1.0",
     "pyarrow>=21.0.0",
     "polars>=1.35.2",
     "numpy>=2.2.6",
@@ -24,30 +23,37 @@ dependencies = [
 
 [project.optional-dependencies]
 serving = [
-    "litserve>=0.2.13",
+    "fastapi>=0.124.0",
+    "orjson>=3.10.0",
     "pydantic-settings>=2.10.1",
+    "uvicorn>=0.38.0",
 ]
 text = [
     "transformers>=4.55.0",
 ]
 docs = [
-    "litserve>=0.2.13",
+    "fastapi>=0.124.0",
     "mkdocs-material>=9.6",
     "mkdocs-jupyter>=0.26.3",
     "mkdocstrings[python]>=0.27",
+    "orjson>=3.10.0",
     "pydantic-settings>=2.10.1",
+    "uvicorn>=0.38.0",
 ]
 
 [dependency-groups]
 dev = [
     "ruff>=0.12.12",
     "pytest>=8.4.1",
+    "pytest-xdist>=3.8.0",
     "ipython>=9.9.0",
     "ipykernel>=6.29.5",
     "nbclient>=0.10.2",
     "nbformat>=5.10.4",
-    "litserve>=0.2.13",
+    "fastapi>=0.124.0",
+    "orjson>=3.10.0",
     "pydantic-settings>=2.10.1",
+    "uvicorn>=0.38.0",
     "ty>=0.0.1a20",
     "pre-commit>=4.3.0",
 ]
@@ -66,6 +72,7 @@ include = ["json2vec*"]
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = ["test_*.py"]
+addopts = ["-n", "auto"]
 
 [tool.ruff]
 line-length = 120

{json2vec-0.4.7 → json2vec-0.4.9}/src/json2vec/__init__.py

@@ -8,15 +8,16 @@ mutation predicates, and the `@preprocess` decorator.
 
 from typing import TYPE_CHECKING, Any
 
+from json2vec import helpers as helpers
+from json2vec.architecture.checkpoint import RollbackCheckpoint
+from json2vec.architecture.mutations import MutationLockCallback, RuntimePlacementCallback
 from json2vec.architecture.root import (
     Model,
-    MutationLockCallback,
     OptimizerConfig,
-    RollbackCheckpoint,
-    RuntimePlacementCallback,
     SchedulerConfig,
 )
 from json2vec.data.datasets import CustomDataModule, PolarsDataModule, StreamingDataModule
+from json2vec.data.processing import MASK_LITERAL, MaskLiteral
 from json2vec.inference.callback import Postprocessor, Writer
 from json2vec.preprocessors import PREPROCESSORS, Preprocessor, PreprocessorMode, preprocess
 from json2vec.structs.enums import (
@@ -38,7 +39,7 @@ from json2vec.structs.experiment import (
     predicate,
     where,
 )
-from json2vec.structs.structure import Array
+from json2vec.structs.structure import Array, Mask
 from json2vec.structs.tree import Address, Leaf
 from json2vec.tensorfields import TENSORFIELDS, DecoderBase, EmbedderBase, Plugin, RequestBase, TensorFieldBase
 from json2vec.tensorfields.extensions.category import Request as Category
@@ -52,23 +53,19 @@ from json2vec.tensorfields.shared.vocabulary import VocabularySyncCallback
 
 if TYPE_CHECKING:
     from json2vec.inference.deployment import (
-        API,
         Accelerator,
-        BatchItem,
         Deployment,
-        ErrorItem,
         Input,
+        JSONBackend,
         ModelSource,
         UpdateOperation,
     )
 
 _SERVING_EXPORTS = {
-    "API",
     "Accelerator",
-    "BatchItem",
     "Deployment",
-    "ErrorItem",
     "Input",
+    "JSONBackend",
     "ModelSource",
     "UpdateOperation",
 }
@@ -81,7 +78,7 @@ def __getattr__(name: str) -> Any:
     try:
         from json2vec.inference import deployment
     except ModuleNotFoundError as error:
-        if error.name in {"litserve", "pydantic_settings"}:
+        if error.name in {"fastapi", "orjson", "pydantic_settings", "uvicorn"}:
             raise ModuleNotFoundError(
                 f"json2vec.{name} requires the serving extra; install with `pip install json2vec[serving]`."
             ) from error
@@ -98,11 +95,9 @@ def __dir__() -> list[str]:
 
 __all__ = [
     "Address",
-    "API",
     "Accelerator",
     "Array",
     "AttentionMode",
-    "BatchItem",
     "Category",
     "Component",
     "CustomDataModule",
@@ -111,11 +106,15 @@ __all__ = [
     "Deployment",
     "EmbedderBase",
     "Entity",
-    "ErrorItem",
+    "helpers",
     "Hyperparameters",
     "Input",
+    "JSONBackend",
     "Leaf",
     "Metric",
+    "MASK_LITERAL",
+    "Mask",
+    "MaskLiteral",
     "Model",
     "ModelSource",
     "MutationLockCallback",

{json2vec-0.4.7 → json2vec-0.4.9}/src/json2vec/architecture/checkpoint.py

@@ -5,7 +5,9 @@ from __future__ import annotations
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import lightning.pytorch as lit
 import torch
+from lightning.pytorch.callbacks import ModelCheckpoint
 from loguru import logger
 
 from json2vec.architecture.graph import ModelGraph
@@ -15,6 +17,46 @@ if TYPE_CHECKING:
     from json2vec.architecture.root import Model
 
 
+class RollbackCheckpoint(ModelCheckpoint):
+    """Checkpoint the best model during fit and restore it into the module at fit end."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        if self.save_weights_only:
+            raise ValueError("RollbackCheckpoint requires full checkpoints; set save_weights_only=False")
+        if self.save_top_k == 0:
+            raise ValueError("RollbackCheckpoint requires at least one saved checkpoint; set save_top_k != 0")
+
+    def on_fit_end(self, trainer: lit.Trainer, pl_module: lit.LightningModule) -> None:
+        from json2vec.architecture.root import Model
+
+        super().on_fit_end(trainer=trainer, pl_module=pl_module)
+        if not isinstance(pl_module, Model):
+            raise TypeError("RollbackCheckpoint can only restore json2vec Model instances")
+
+        best_model_path = self.best_model_path
+        if not best_model_path:
+            raise RuntimeError("RollbackCheckpoint did not find a best checkpoint to restore")
+
+        strategy = getattr(trainer, "strategy", None)
+        if strategy is not None:
+            strategy.barrier("rollback_checkpoint_load")
+            checkpoint = strategy.checkpoint_io.load_checkpoint(
+                best_model_path,
+                map_location=pl_module.device,
+                weights_only=False,
+            )
+        else:
+            checkpoint = torch.load(best_model_path, weights_only=False, map_location=pl_module.device)
+
+        pl_module.restore_checkpoint_state(checkpoint)
+        logger.bind(
+            component="checkpoint",
+            checkpoint=best_model_path,
+            score=self.best_model_score,
+        ).info("rolled back Model to best checkpoint")
+
+
 class CheckpointState:
     """Save, load, and restore model state without owning the public facade."""
 

{json2vec-0.4.7 → json2vec-0.4.9}/src/json2vec/architecture/contracts.py

@@ -3,9 +3,9 @@
 from __future__ import annotations
 
 from collections.abc import Iterator, Mapping
-from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
+import pydantic
 import torch
 from tensordict import TensorDict
 
@@ -34,12 +34,13 @@ ContractSignature = tuple[Any, ...]
 ContractScope = tuple[str, int, int, ContractSignature]
 
 
-@dataclass
-class ContractScheduler:
+class ContractScheduler(pydantic.BaseModel):
     """Deterministic backoff scheduler for expensive forward contract checks."""
 
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
     periodic_interval: int = 1024
-    _counts: dict[ContractScope, int] = field(default_factory=dict)
+    _counts: dict[ContractScope, int] = pydantic.PrivateAttr(default_factory=dict)
 
     def reset(self) -> None:
         self._counts.clear()

{json2vec-0.4.7 → json2vec-0.4.9}/src/json2vec/architecture/mutations.py

@@ -4,12 +4,17 @@ from __future__ import annotations
 
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from dataclasses import dataclass
+from functools import partialmethod, wraps
 from typing import TYPE_CHECKING, Any
 
+import lightning.pytorch as lit
+import pydantic
+import torch
+from lightning.pytorch import Callback
 from loguru import logger
 
 from json2vec.architecture.graph import ModelGraph
+from json2vec.structs.enums import Strata
 from json2vec.structs.experiment import NodeAttribute, NodePredicate, SchemaField
 from json2vec.structs.structure import Array
 from json2vec.structs.tree import Leaf, Node
@@ -20,12 +25,85 @@ if TYPE_CHECKING:
 _MISSING = object()
 
 
-@dataclass(frozen=True)
-class AttributeChange:
+def immutable(name: str | Strata) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    def decorator(method: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(method)
+        def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
+            locks = self.locks
+            locks[name] += 1
+            try:
+                return method(self, *args, **kwargs)
+            finally:
+                if locks[name] <= 1:
+                    locks.pop(name, None)
+                else:
+                    locks[name] -= 1
+
+        return wrapped
+
+    return decorator
+
+
+class MutationLockCallback(Callback):
+    """Prevent runtime schema mutations while Lightning owns an active loop."""
+
+    locks: tuple[Strata, ...] = (Strata.train, Strata.validate, Strata.test, Strata.predict)
+
+    def _on_loop_start(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
+        pl_module.locks[strata] += 1
+
+    def _on_loop_end(self, trainer: lit.Trainer, pl_module: "Model", strata: Strata) -> None:
+        locks = pl_module.locks
+        if locks[strata] <= 1:
+            locks.pop(strata, None)
+        else:
+            locks[strata] -= 1
+
+    def on_exception(
+        self,
+        trainer: lit.Trainer,
+        pl_module: "Model",
+        exception: BaseException,
+    ) -> None:  # ty:ignore[invalid-method-override]
+        for lock in self.locks:
+            pl_module.locks.pop(lock, None)
+
+    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
+    on_train_end = partialmethod(_on_loop_end, strata=Strata.train)
+    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
+    on_validation_end = partialmethod(_on_loop_end, strata=Strata.validate)
+    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
+    on_test_end = partialmethod(_on_loop_end, strata=Strata.test)
+    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
+    on_predict_end = partialmethod(_on_loop_end, strata=Strata.predict)
+
+
+class RuntimePlacementCallback(Callback):
+    """Move late-created modules onto the Lightning module's active device."""
+
+    def _on_loop_start(self, trainer: lit.Trainer, pl_module: lit.LightningModule, strata: Strata) -> None:
+        device = getattr(pl_module, "device", None)
+        if isinstance(device, torch.device):
+            pl_module.to(device=device)
+
+    on_train_start = partialmethod(_on_loop_start, strata=Strata.train)
+    on_validation_start = partialmethod(_on_loop_start, strata=Strata.validate)
+    on_test_start = partialmethod(_on_loop_start, strata=Strata.test)
+    on_predict_start = partialmethod(_on_loop_start, strata=Strata.predict)
+
+
+class AttributeChange(pydantic.BaseModel):
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
     node: Node
     name: str
     original: Any
     definition_attribute: bool
+    address: str
+    node_name: str
+    node_type: str
+    changed: Any = _MISSING
+    changed_address: Any = _MISSING
 
 
 class SchemaEditor:
@@ -34,6 +112,12 @@ class SchemaEditor:
     def __init__(self, module: "Model") -> None:
         self.module = module
 
+    def _assert_mutation_allowed(self, action: str) -> None:
+        active = tuple(name for name, count in self.module.locks.items() if count > 0)
+        if active:
+            labels = ", ".join(active)
+            raise RuntimeError(f"model.{action}(...) cannot run while the model is in an active loop: {labels}")
+
     def select(
         self,
         *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
@@ -56,7 +140,7 @@ class SchemaEditor:
         use_cache: bool = False,
         **values: Any,
     ) -> None:
-        self.module._assert_mutation_allowed("update")
+        self._assert_mutation_allowed("update")
         values = self.module.hyperparameters.update_values(values)
         changes = self._attribute_changes(
             values=values,
@@ -84,7 +168,7 @@ class SchemaEditor:
         include_root: bool = True,
         use_cache: bool = True,
     ) -> None:
-        self.module._assert_mutation_allowed("extend")
+        self._assert_mutation_allowed("extend")
         parent, field_count = self._extend_target(*args, include_root=include_root, use_cache=use_cache)
         self.module.hyperparameters.extend(*args, include_root=include_root, use_cache=use_cache)
         ModelGraph.rebuild(self.module)
@@ -103,7 +187,7 @@ class SchemaEditor:
         include_root: bool = False,
         use_cache: bool = True,
     ) -> None:
-        self.module._assert_mutation_allowed("delete")
+        self._assert_mutation_allowed("delete")
         roots = self._delete_roots(*predicates, include_root=include_root, use_cache=use_cache)
         self.module.hyperparameters.delete(*predicates, include_root=include_root, use_cache=use_cache)
         ModelGraph.rebuild(self.module)
@@ -123,7 +207,7 @@ class SchemaEditor:
         use_cache: bool = True,
         descendants: bool = False,
     ) -> None:
-        self.module._assert_mutation_allowed("reset")
+        self._assert_mutation_allowed("reset")
         selected = self.module.hyperparameters.select(
             *predicates,
             include_root=include_root,
@@ -154,7 +238,7 @@ class SchemaEditor:
         use_cache: bool = False,
         **values: Any,
     ) -> Iterator[None]:
-        self.module._assert_mutation_allowed("override")
+        self._assert_mutation_allowed("override")
         values = self.module.hyperparameters.update_values(values)
         changes = self._attribute_changes(
             values=values,
@@ -208,6 +292,9 @@ class SchemaEditor:
                         name=name,
                         original=getattr(node, name, _MISSING),
                         definition_attribute=_is_definition_attribute(node, name),
+                        address=str(node.address),
+                        node_name=node.name,
+                        node_type=node.type,
                     )
                 )
 
@@ -280,29 +367,55 @@ class SchemaEditor:
 
     def _log_attribute_changes(self, action: str, changes: list[AttributeChange], *, restored: bool = False) -> None:
         for change in changes:
+            current_address = str(change.node.address)
             value = change.original if restored else getattr(change.node, change.name, _MISSING)
+            if not restored:
+                change.changed = value
+                change.changed_address = current_address
+            previous_value = change.changed if restored else change.original
+            previous_address = change.changed_address if restored else change.address
+            if previous_address is _MISSING:
+                previous_address = change.address
+            address_context = (
+                current_address if previous_address == current_address else f"{previous_address} -> {current_address}"
+            )
+            value_text = _format_log_value(value)
+            previous_value_text = _format_log_value(previous_value)
             logger.bind(
                 component="schema_mutation",
                 action=action,
-                address=str(change.node.address),
-                node_type=change.node.type,
+                address=current_address,
+                previous_address=previous_address,
+                node_name=change.node.name,
+                previous_node_name=change.node_name,
+                node_type=change.node_type,
                 attribute=change.name,
                 definition_attribute=change.definition_attribute,
-                value=_format_log_value(value),
-                previous_value=_format_log_value(change.original),
-            ).info("restored schema node attribute" if restored else "mutated schema node attribute")
+                value=value_text,
+                previous_value=previous_value_text,
+                change=f"{change.name}: {previous_value_text} -> {value_text}",
+            ).info(
+                "{} {}: {} {} -> {}",
+                "restored" if restored else "mutated",
+                address_context,
+                change.name,
+                previous_value_text,
+                value_text,
+            )
 
     def _log_node_mutation(self, *, action: str, message: str, node: Node, **kwargs: Any) -> None:
         extra = {key: str(value.address) if isinstance(value, Node) else value for key, value in kwargs.items()}
+        context = _format_node_log_context(node, extra)
         logger.bind(
             component="schema_mutation",
             action=action,
             address=str(node.address),
             node_type=node.type,
+            node_name=node.name,
             attribute=None,
             definition_attribute=None,
             **extra,
-        ).info(message)
+        ).info("{} {}", message, context)
 
 
 def _has_node_attribute(node: Node, name: str) -> bool:
@@ -321,3 +434,13 @@ def _format_log_value(value: Any) -> str:
 
     text = repr(value)
     return text if len(text) <= 160 else f"{text[:157]}..."
+
+
+def _format_node_log_context(node: Node, extra: dict[str, Any]) -> str:
+    parts = [str(node.address)]
+    if parent := extra.get("parent"):
+        parts.append(f"under {parent}")
+    if "descendants" in extra:
+        parts.append(f"descendants={extra['descendants']}")
+
+    return " ".join(parts)