json2vec 0.4.7__tar.gz → 0.4.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.7
+Version: 0.4.8
 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.8}/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.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "json2vec"
-version = "0.4.7"
+version = "0.4.8"
 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.8}/src/json2vec/__init__.py

@@ -52,23 +52,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 +77,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 +94,9 @@ def __dir__() -> list[str]:
 
 __all__ = [
     "Address",
-    "API",
     "Accelerator",
     "Array",
     "AttentionMode",
-    "BatchItem",
     "Category",
     "Component",
     "CustomDataModule",
@@ -111,9 +105,9 @@ __all__ = [
     "Deployment",
     "EmbedderBase",
     "Entity",
-    "ErrorItem",
     "Hyperparameters",
     "Input",
+    "JSONBackend",
     "Leaf",
     "Metric",
     "Model",

{json2vec-0.4.7 → json2vec-0.4.8}/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.8}/src/json2vec/architecture/mutations.py

@@ -4,9 +4,9 @@ from __future__ import annotations
 
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+import pydantic
 from loguru import logger
 
 from json2vec.architecture.graph import ModelGraph
@@ -20,12 +20,18 @@ if TYPE_CHECKING:
 _MISSING = object()
 
 
-@dataclass(frozen=True)
-class AttributeChange:
+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:
@@ -208,6 +214,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 +289,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 +356,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)

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/architecture/root.py

@@ -13,15 +13,16 @@ from beartype import beartype
 from lightning.pytorch import Callback
 from lightning.pytorch.callbacks import ModelCheckpoint
 from loguru import logger
+from rich.text import Text
 from tensordict import TensorDict
 
 from json2vec.architecture.checkpoint import CheckpointState
 from json2vec.architecture.contracts import ContractScheduler
 from json2vec.architecture.graph import ModelGraph
 from json2vec.architecture.mutations import SchemaEditor
-from json2vec.architecture.plot import PlotMode
 from json2vec.architecture.runtime import ModelRuntime, Postprocessor, Preprocessor, step
 from json2vec.data.datasets.base import EncodedBatch, EncodedInput
+from json2vec.logging.throughput import ThroughputLogger
 from json2vec.structs.enums import AttentionMode, Strata
 from json2vec.structs.experiment import (
     Hyperparameters,
@@ -30,7 +31,7 @@ from json2vec.structs.experiment import (
     SchemaField,
 )
 from json2vec.structs.packages import Prediction
-from json2vec.structs.tree import Address, Node, Rate
+from json2vec.structs.tree import Address, Node, Rate, Renderable
 from json2vec.tensorfields.base import TENSORFIELDS, Plugin, TensorFieldBase
 
 OptimizerConfig = torch.optim.Optimizer | Callable[["Model"], torch.optim.Optimizer]
@@ -137,12 +138,12 @@ class RollbackCheckpoint(ModelCheckpoint):
         ).info("rolled back Model to best checkpoint")
 
 
-class Model(lit.LightningModule):
+class Model(lit.LightningModule, Renderable):
     """Neural model generated from a `json2vec` schema tree.
 
     `Model` owns the schema hyperparameters, tensorfield embedders, array
     encoders, decoders, and convenience methods for prediction, checkpointing,
-    plotting, and schema mutation.
+    schema display and mutation.
 
     Example:
         ```python
@@ -265,6 +266,45 @@ class Model(lit.LightningModule):
         self._contract_generation += 1
         self._contract_scheduler.reset()
 
+    def __rich_console__(self, console, options):
+        parameters = sum(parameter.numel() for parameter in self.parameters())
+        heading = Text()
+        heading.append(type(self).__name__, style=self.RICH_NAME_STYLE)
+        heading.append(" ")
+        heading.append("[model]", style=self.RICH_TYPE_STYLE)
+        for name, value in (
+            ("batch_size", self.batch_size),
+            ("d_model", self.hyperparameters.d_model),
+            ("parameters", f"{parameters:,}"),
+            ("arrays", len(self.hyperparameters.arrays)),
+            ("fields", len(self.hyperparameters.active_requests)),
+            ("targets", len(self.hyperparameters.target)),
+            ("embeds", len(self.hyperparameters.embed)),
+        ):
+            heading.append(" ")
+            heading.append(f"{name}=", style="dim")
+            heading.append(str(value), style="cyan")
+        yield heading
+
+        lines = list(self.hyperparameters.fields.__rich_console__(console, options))
+        if not lines:
+            return
+        first = Text()
+        first.append("`-- ", style=self.RICH_TREE_STYLE)
+        if isinstance(lines[0], Text):
+            first.append_text(lines[0])
+        else:
+            first.append(str(lines[0]))
+        yield first
+        for line in lines[1:]:
+            nested = Text()
+            nested.append("    ", style=self.RICH_TREE_STYLE)
+            if isinstance(line, Text):
+                nested.append_text(line)
+            else:
+                nested.append(str(line))
+            yield nested
+
     def select(
         self,
         *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
@@ -382,6 +422,8 @@ class Model(lit.LightningModule):
             callbacks.append(RuntimePlacementCallback())
         if MutationLockCallback not in attached_callback_types:
             callbacks.append(MutationLockCallback())
+        if ThroughputLogger not in attached_callback_types:
+            callbacks.append(ThroughputLogger())
 
         for request in self.hyperparameters.active_requests.values():
             plugin: Plugin = TENSORFIELDS[request.type]
@@ -438,26 +480,6 @@ class Model(lit.LightningModule):
             if hasattr(node, "embedder") and hasattr(node.embedder, "interprocess_encoding_context")
         }
 
-    def plot(
-        self,
-        address: Address | str | None = None,
-        detail: bool = False,
-        out: str | Path | None = None,
-        mode: PlotMode = "schema",
-    ) -> None:
-        """Print a Rich model visualization.
-
-        Args:
-            address: Optional subtree address to render.
-            detail: Include tensorfield-specific detail sections.
-            out: Optional output path for the rendered console text.
-            mode: Plot mode. Supported values are `schema`, `state`, `flow`,
-                and `debug`.
-        """
-        from json2vec.architecture.plot import plot
-
-        return plot(module=self, address=address, detail=detail, out=out, mode=mode)
-
     @beartype
     def save(self, pathname: str | Path) -> str | Path:
         """Save model weights and schema hyperparameters to a checkpoint."""

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/data/datasets/base.py

@@ -90,5 +90,13 @@ def _is_assigned_to_worker(shard_key: str, worker_id: int, num_workers: int) ->
     return owner == worker_id
 
 
+def share_interprocess_encoding_context(context: InterprocessEncodingContext) -> None:
+    """Opt encoding context resources into multiprocessing-safe storage."""
+    for field_context in context.values():
+        share = getattr(field_context, "share", None)
+        if callable(share):
+            share()
+
+
 def identity(data: Any) -> Any:
     return data

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/data/datasets/custom.py

@@ -22,6 +22,7 @@ from json2vec.data.datasets.base import (
     SampleRate,
     StrataMap,
     identity,
+    share_interprocess_encoding_context,
 )
 from json2vec.data.iterables import (
     JMESPathResolutionMonitor,
@@ -299,15 +300,23 @@ class CustomDataModule(lit.LightningDataModule):
                 return None
             raise ValueError(f"no dataset configured for strata: {strata}")
 
+        workers = self.num_workers[strata]
+        if workers is None:
+            workers = os.cpu_count() or 0
+
+        interprocess_encoding_context = self.interprocess_encoding_context
+        if strata == Strata.train and workers > 0:
+            share_interprocess_encoding_context(interprocess_encoding_context)
+
         return custom_dataloader(
             hyperparameters=self.hyperparameters,
             dataset=self.datasets[strata],
             preprocessor=self.preprocessor,
             preprocessor_kwargs=self.preprocessor_kwargs,
-            interprocess_encoding_context=self.interprocess_encoding_context,
+            interprocess_encoding_context=interprocess_encoding_context,
             batch_size=self.batch_size,
             strata=strata,
-            num_workers=self.num_workers[strata],
+            num_workers=workers,
             persistent_workers=self.persistent_workers[strata],
             pin_memory=self.pin_memory[strata],
             observation_buffer_size=self.observation_buffer_size[strata],

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/data/datasets/polars.py

@@ -26,6 +26,7 @@ from json2vec.data.datasets.base import (
     _is_assigned_to_worker,
     _worker_identity,
     identity,
+    share_interprocess_encoding_context,
 )
 from json2vec.data.iterables import (
     JMESPathResolutionMonitor,
@@ -355,15 +356,23 @@ class PolarsDataModule(lit.LightningDataModule):
                 return None
             raise ValueError(f"no dataframe configured for strata: {strata}")
 
+        workers = self.num_workers[strata]
+        if workers is None:
+            workers = os.cpu_count() or 0
+
+        interprocess_encoding_context = self.interprocess_encoding_context
+        if strata == Strata.train and workers > 0:
+            share_interprocess_encoding_context(interprocess_encoding_context)
+
         return polars_dataloader(
             hyperparameters=self.hyperparameters,
             dataframe=self.dataframes[strata],
             preprocessor=self.preprocessor,
             preprocessor_kwargs=self.preprocessor_kwargs,
-            interprocess_encoding_context=self.interprocess_encoding_context,
+            interprocess_encoding_context=interprocess_encoding_context,
             batch_size=self.batch_size,
             strata=strata,
-            num_workers=self.num_workers[strata],
+            num_workers=workers,
             persistent_workers=self.persistent_workers[strata],
             pin_memory=self.pin_memory[strata],
             sharding=self.sharding[strata],

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/data/datasets/streaming.py

@@ -31,6 +31,7 @@ from json2vec.data.datasets.base import (
     _is_assigned_to_worker,
     _worker_identity,
     identity,
+    share_interprocess_encoding_context,
 )
 from json2vec.data.iterables import (
     JMESPathResolutionMonitor,
@@ -500,6 +501,14 @@ class StreamingDataModule(lit.LightningDataModule):
         global_rank = getattr(trainer, "global_rank", None)
         world_size = getattr(trainer, "world_size", None)
 
+        workers = self.num_workers[strata]
+        if workers is None:
+            workers = os.cpu_count() or 0
+
+        interprocess_encoding_context = self.interprocess_encoding_context
+        if strata == Strata.train and workers > 0:
+            share_interprocess_encoding_context(interprocess_encoding_context)
+
         return dataloader(
             hyperparameters=self.hyperparameters,
             root=self.root,
@@ -507,10 +516,10 @@ class StreamingDataModule(lit.LightningDataModule):
             pattern=pattern,
             preprocessor=self.preprocessor,
             preprocessor_kwargs=self.preprocessor_kwargs,
-            interprocess_encoding_context=self.interprocess_encoding_context,
+            interprocess_encoding_context=interprocess_encoding_context,
             batch_size=self.batch_size,
             strata=strata,
-            num_workers=self.num_workers[strata],
+            num_workers=workers,
             persistent_workers=self.persistent_workers[strata],
             pin_memory=self.pin_memory[strata],
             sharding=self.sharding[strata],

{json2vec-0.4.7 → json2vec-0.4.8}/src/json2vec/data/iterables.py

@@ -4,8 +4,9 @@ from __future__ import annotations
 
 import inspect
 import random
+import re
 from collections import Counter
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
 from functools import cache
 from typing import Annotated, Any, TypeVar, cast
 
@@ -138,6 +139,56 @@ def query(expression: str) -> jmespath.parser.ParsedResult:
     return jmespath.compile(expression=f"[*]{expression}")
 
 
+@cache
+def compile_query_extractor(expression: str) -> Callable[[Any], Any] | None:
+    """Compile a direct extractor for simple JMESPath projection queries."""
+    if not expression.startswith("[*]"):
+        return None
+
+    operations: list[tuple[str, str | None]] = [("projection", None), ("projection", None)]
+    index = 3
+    while index < len(expression):
+        if expression.startswith("[*]", index):
+            operations.append(("projection", None))
+            index += 3
+            continue
+
+        if not expression.startswith(".", index):
+            return None
+
+        match = re.match(r"\.([A-Za-z_][A-Za-z0-9_]*)", expression[index:])
+        if match is None:
+            return None
+
+        operations.append(("field", match.group(1)))
+        index += len(match.group(0))
+
+    def search(value: Any) -> Any:
+        def apply(item: Any, operation_index: int) -> Any:
+            if operation_index == len(operations):
+                return item
+
+            operation, key = operations[operation_index]
+            if operation == "field":
+                if not isinstance(item, dict) or key not in item:
+                    return None
+                return apply(item[key], operation_index + 1)
+
+            if not isinstance(item, list):
+                return None
+
+            out = []
+            for child in item:
+                result = apply(child, operation_index + 1)
+                if result is not None:
+                    out.append(result)
+            return out
+
+        return apply(value, 0)
+
+    return search
+
+
 class JMESPathResolutionMonitor(pydantic.BaseModel):
     every: Annotated[int, pydantic.Field(gt=0)] = 1000
 
@@ -192,9 +243,11 @@ def encode(
         if expression is None:
             raise ValueError(f"request '{address}' must define query")
 
-        # `request.query` is relative to a processed observation. `query(...)`
-        # adds the outer batch selector before JMESPath searches `batch`.
-        result = query(expression).search(batch)
+        # `request.query` is relative to a processed observation. Direct
+        # extractors add the outer batch selector; `query(...)` does the same
+        # before JMESPath searches `batch`.
+        extractor = compile_query_extractor(expression)
+        result = extractor(batch) if extractor is not None else query(expression).search(batch)
         if jmespath_resolution_monitor is not None:
             jmespath_resolution_monitor.observe(address=address, expression=expression, result=result)