json2vec 0.4.4__tar.gz → 0.4.5__tar.gz

{json2vec-0.4.4/src/json2vec.egg-info → json2vec-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.4
+Version: 0.4.5
 Summary: {...} -> [*]
 License-Expression: Apache-2.0
 Requires-Python: >=3.12

{json2vec-0.4.4 → json2vec-0.4.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "json2vec"
-version = "0.4.4"
+version = "0.4.5"
 description = "{...} -> [*]"
 readme = "README.md"
 license = "Apache-2.0"

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/architecture/contracts.py

@@ -93,7 +93,7 @@ def sanitize(
         require_core_tensors(module, address, tensorfield)
         require_tensor_devices(module, address, tensorfield)
         require_target_contract(module, address, tensorfield, strata=normalized)
-        require_mask_contract(module, address, tensorfield)
+        require_mask_contract(module, address, tensorfield, strata=normalized)
 
 
 def is_backoff_index(index: int, *, periodic_interval: int) -> bool:
@@ -292,7 +292,7 @@ def require_tensor_devices(module: "Model", address: Address, tensorfield: Tenso
         )
 
 
-def require_mask_contract(module: "Model", address: Address, tensorfield: TensorFieldBase) -> None:
+def require_mask_contract(module: "Model", address: Address, tensorfield: TensorFieldBase, *, strata: Strata) -> None:
     state = tensorfield.state
     trainable = tensorfield.trainable
     is_masked = state.eq(Tokens.masked.value)
@@ -301,7 +301,7 @@ def require_mask_contract(module: "Model", address: Address, tensorfield: Tensor
     if trainable.any() and not state.masked_select(trainable).eq(Tokens.masked.value).all():
         raise ForwardContractError(f"forward input '{address}' trainable positions must have masked state")
 
-    if not is_target and (is_masked & ~trainable).any():
+    if strata != Strata.predict and not is_target and (is_masked & ~trainable).any():
         raise ForwardContractError(f"forward input '{address}' has masked state where trainable is false")
 
     if not trainable.any():

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/architecture/graph.py

@@ -8,6 +8,8 @@ from typing import TYPE_CHECKING
 import torch
 
 from json2vec.architecture.node import NodeModule
+from json2vec.data.datasets.base import EncodedInput
+from json2vec.structs.enums import Strata
 from json2vec.structs.experiment import Hyperparameters
 from json2vec.structs.tree import Address, Node
 
@@ -19,9 +21,19 @@ class ModelGraph:
     """Build and rebuild runtime modules from schema hyperparameters."""
 
     @staticmethod
-    def build(hyperparameters: Hyperparameters, batch_size: int) -> tuple[torch.nn.ModuleDict, torch.Tensor]:
+    def example_forward_kwargs(hyperparameters: Hyperparameters, batch_size: int) -> dict[str, EncodedInput | Strata]:
         from json2vec.data.iterables import mock
 
+        return {
+            "inputs": mock(hyperparameters=hyperparameters, batch_size=batch_size),
+            "strata": Strata.predict,
+        }
+
+    @staticmethod
+    def build(
+        hyperparameters: Hyperparameters,
+        batch_size: int,
+    ) -> tuple[torch.nn.ModuleDict, dict[str, EncodedInput | Strata]]:
         nodes: torch.nn.ModuleDict[str, NodeModule] = torch.nn.ModuleDict()
 
         for address in hyperparameters.requests | hyperparameters.arrays:
@@ -31,7 +43,7 @@ class ModelGraph:
                 batch_size=batch_size,
             )
 
-        return nodes, mock(hyperparameters=hyperparameters, batch_size=batch_size)
+        return nodes, ModelGraph.example_forward_kwargs(hyperparameters=hyperparameters, batch_size=batch_size)
 
     @staticmethod
     def install(module: "Model") -> None:
@@ -72,8 +84,6 @@ class ModelGraph:
 
     @staticmethod
     def reset_selected(module: "Model", selected: list[Node], *, descendants: bool = False) -> None:
-        from json2vec.data.iterables import mock
-
         selected_by_address: dict[Address, Node] = {}
         for node in selected:
             if node.address in module.nodes:
@@ -94,7 +104,10 @@ class ModelGraph:
                 batch_size=module.batch_size,
             )
 
-        module.example_input_array = mock(hyperparameters=module.hyperparameters, batch_size=module.batch_size)
+        module.example_input_array = ModelGraph.example_forward_kwargs(
+            hyperparameters=module.hyperparameters,
+            batch_size=module.batch_size,
+        )
         device = module.device
         if isinstance(device, torch.device):
             module.to(device=device)

json2vec-0.4.5/src/json2vec/architecture/mutations.py

@@ -0,0 +1,323 @@
+"""Model-facing schema mutation orchestration."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+from json2vec.architecture.graph import ModelGraph
+from json2vec.structs.experiment import NodeAttribute, NodePredicate, SchemaField
+from json2vec.structs.structure import Array
+from json2vec.structs.tree import Leaf, Node
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import Model
+
+_MISSING = object()
+
+
+@dataclass(frozen=True)
+class AttributeChange:
+    node: Node
+    name: str
+    original: Any
+    definition_attribute: bool
+
+
+class SchemaEditor:
+    """Coordinate schema mutations with runtime graph rebuilds."""
+
+    def __init__(self, module: "Model") -> None:
+        self.module = module
+
+    def select(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        include_root: bool = True,
+        use_cache: bool = True,
+    ) -> list[Node]:
+        return self.module.hyperparameters.select(
+            *predicates,
+            include_root=include_root,
+            use_cache=use_cache,
+        )
+
+    def update(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        strict: bool = True,
+        allow_extra: bool = False,
+        include_root: bool = True,
+        validate: bool = True,
+        use_cache: bool = False,
+        **values: Any,
+    ) -> None:
+        self.module._assert_mutation_allowed("update")
+        values = self.module.hyperparameters.update_values(values)
+        changes = self._attribute_changes(
+            values=values,
+            predicates=predicates,
+            allow_extra=allow_extra,
+            include_root=include_root,
+            use_cache=use_cache,
+        )
+        self.module.hyperparameters.update(
+            *predicates,
+            strict=strict,
+            allow_extra=allow_extra,
+            include_root=include_root,
+            validate=validate,
+            use_cache=use_cache,
+            **values,
+        )
+        ModelGraph.rebuild(self.module)
+        self.module._reset_contracts()
+        self._log_attribute_changes("update", changes)
+
+    def extend(
+        self,
+        *args: NodePredicate | NodeAttribute | Callable[[Node], bool] | SchemaField,
+        include_root: bool = True,
+        use_cache: bool = True,
+    ) -> None:
+        self.module._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)
+        self.module._reset_contracts()
+        for field in parent.fields[-field_count:]:
+            self._log_node_mutation(
+                action="extend",
+                message="extended schema node",
+                node=field,
+                parent=parent,
+            )
+
+    def delete(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        include_root: bool = False,
+        use_cache: bool = True,
+    ) -> None:
+        self.module._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)
+        self.module._reset_contracts()
+        for node in roots:
+            self._log_node_mutation(
+                action="delete",
+                message="deleted schema node",
+                node=node,
+                descendants=len(getattr(node, "descendants", ())),
+            )
+
+    def reset(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        include_root: bool = True,
+        use_cache: bool = True,
+        descendants: bool = False,
+    ) -> None:
+        self.module._assert_mutation_allowed("reset")
+        selected = self.module.hyperparameters.select(
+            *predicates,
+            include_root=include_root,
+            use_cache=use_cache,
+        )
+        if not selected:
+            raise ValueError("reset matched no nodes")
+
+        nodes = self._runtime_reset_nodes(selected, descendants=descendants)
+        ModelGraph.reset_selected(self.module, selected, descendants=descendants)
+        self.module._reset_contracts()
+        for node in nodes:
+            self._log_node_mutation(
+                action="reset",
+                message="reset runtime node",
+                node=node,
+                descendants=descendants,
+            )
+
+    @contextmanager
+    def override(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        strict: bool = True,
+        allow_extra: bool = False,
+        include_root: bool = True,
+        validate: bool = True,
+        use_cache: bool = False,
+        **values: Any,
+    ) -> Iterator[None]:
+        self.module._assert_mutation_allowed("override")
+        values = self.module.hyperparameters.update_values(values)
+        changes = self._attribute_changes(
+            values=values,
+            predicates=predicates,
+            allow_extra=allow_extra,
+            include_root=include_root,
+            use_cache=use_cache,
+        )
+        entered = False
+        try:
+            with self.module.hyperparameters.override(
+                *predicates,
+                strict=strict,
+                allow_extra=allow_extra,
+                include_root=include_root,
+                validate=validate,
+                use_cache=use_cache,
+                **values,
+            ):
+                entered = True
+                ModelGraph.rebuild(self.module)
+                self.module._reset_contracts()
+                self._log_attribute_changes("override", changes)
+                yield
+        finally:
+            ModelGraph.rebuild(self.module)
+            self.module._reset_contracts()
+            if entered:
+                self._log_attribute_changes("override_restore", changes, restored=True)
+
+    def _attribute_changes(
+        self,
+        *,
+        values: dict[str, Any],
+        predicates: tuple[NodePredicate | NodeAttribute | Callable[[Node], bool], ...],
+        allow_extra: bool,
+        include_root: bool,
+        use_cache: bool,
+    ) -> list[AttributeChange]:
+        nodes = self.module.hyperparameters.select(*predicates, include_root=include_root, use_cache=use_cache)
+        changes: list[AttributeChange] = []
+        for node in nodes:
+            can_apply_extra = allow_extra and getattr(type(node), "model_config", {}).get("extra") == "allow"
+            for name in values:
+                if not (_has_node_attribute(node, name) or can_apply_extra):
+                    continue
+
+                changes.append(
+                    AttributeChange(
+                        node=node,
+                        name=name,
+                        original=getattr(node, name, _MISSING),
+                        definition_attribute=_is_definition_attribute(node, name),
+                    )
+                )
+
+        return changes
+
+    def _extend_target(
+        self,
+        *args: NodePredicate | NodeAttribute | Callable[[Node], bool] | SchemaField,
+        include_root: bool,
+        use_cache: bool,
+    ) -> tuple[Array, int]:
+        predicates: list[NodePredicate | NodeAttribute | Callable[[Node], bool]] = []
+        field_count = 0
+        reading_fields = False
+
+        for item in args:
+            if isinstance(item, (Array, Leaf)):
+                reading_fields = True
+                field_count += 1
+                continue
+
+            if reading_fields:
+                raise TypeError("extend predicates must come before new schema fields")
+
+            predicates.append(item)
+
+        if field_count == 0:
+            raise ValueError("extend requires at least one schema field")
+
+        candidates = [
+            node
+            for node in self.module.hyperparameters.select(*predicates, include_root=include_root, use_cache=use_cache)
+            if isinstance(node, Array)
+        ]
+        if len(candidates) != 1:
+            raise ValueError(f"extend requires exactly one matching array node, found {len(candidates)}")
+
+        return candidates[0], field_count
+
+    def _delete_roots(
+        self,
+        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
+        include_root: bool,
+        use_cache: bool,
+    ) -> list[Node]:
+        selected = self.module.hyperparameters.select(*predicates, include_root=include_root, use_cache=use_cache)
+        selected_ids = {id(node) for node in selected}
+        return [
+            node
+            for node in selected
+            if not any(
+                id(ancestor) in selected_ids
+                for ancestor in getattr(node, "ancestors", ())
+                if ancestor is not self.module.hyperparameters
+            )
+        ]
+
+    def _runtime_reset_nodes(self, selected: list[Node], *, descendants: bool) -> list[Node]:
+        nodes: dict[str, Node] = {}
+        for node in selected:
+            if node.address in self.module.nodes:
+                nodes[str(node.address)] = node
+
+            if descendants:
+                for descendant in getattr(node, "descendants", ()):
+                    if descendant.address in self.module.nodes:
+                        nodes[str(descendant.address)] = descendant
+
+        return list(nodes.values())
+
+    def _log_attribute_changes(self, action: str, changes: list[AttributeChange], *, restored: bool = False) -> None:
+        for change in changes:
+            value = change.original if restored else getattr(change.node, change.name, _MISSING)
+            logger.bind(
+                component="schema_mutation",
+                action=action,
+                address=str(change.node.address),
+                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")
+
+    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()}
+        logger.bind(
+            component="schema_mutation",
+            action=action,
+            address=str(node.address),
+            node_type=node.type,
+            attribute=None,
+            definition_attribute=None,
+            **extra,
+        ).info(message)
+
+
+def _has_node_attribute(node: Node, name: str) -> bool:
+    fields = getattr(type(node), "model_fields", {})
+    extra = getattr(node, "model_extra", None) or {}
+    return name in fields or name in extra or hasattr(node, name)
+
+
+def _is_definition_attribute(node: Node, name: str) -> bool:
+    return name in getattr(type(node), "model_fields", {})
+
+
+def _format_log_value(value: Any) -> str:
+    if value is _MISSING:
+        return "<missing>"
+
+    text = repr(value)
+    return text if len(text) <= 160 else f"{text[:157]}..."

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/architecture/root.py

@@ -18,10 +18,10 @@ 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 EvaluationResult, ModelRuntime, Postprocessor, PreprocessFn, step
-from json2vec.architecture.schema_editor import SchemaEditor
-from json2vec.data.datasets.base import EncodedBatch
+from json2vec.data.datasets.base import EncodedBatch, EncodedInput
 from json2vec.structs.enums import AttentionMode, Strata
 from json2vec.structs.experiment import (
     Hyperparameters,
@@ -520,24 +520,38 @@ class Model(lit.LightningModule):
     ) -> tuple[dict[Address, dict[str, Any]], dict[Address, dict[str, Any]]]:
         return ModelRuntime.write(self, predictions)
 
+    @immutable("inference")
+    def encode(
+        self,
+        batch: EncodedBatch | list[dict[str, Any]],
+        preprocess: PreprocessFn | None = None,
+        strata: Strata | str = Strata.predict,
+    ) -> EncodedInput:
+        """Return encoded tensorfield inputs for raw or processed observations."""
+        return ModelRuntime.encode(
+            self,
+            batch=batch,
+            preprocess=preprocess,
+            strata=strata,
+        )
+
     @immutable("inference")
     def evaluate(
         self,
         batch: EncodedBatch | list[dict[str, Any]],
         preprocess: PreprocessFn | None = None,
         postprocess: Postprocessor | None = None,
-    ) -> tuple[dict[Address, dict[str, Any]], dict[Address, dict[str, Any]]]:
+    ) -> EvaluationResult:
         """Run prediction and embedding for encoded or raw observations.
 
         If `preprocess` is omitted, raw records are encoded unchanged.
         """
-        result: EvaluationResult = ModelRuntime.evaluate(
+        return ModelRuntime.evaluate(
             self,
             batch=batch,
             preprocess=preprocess,
             postprocess=postprocess,
         )
-        return result.as_tuple()
 
     def predict(
         self,
@@ -546,12 +560,14 @@ class Model(lit.LightningModule):
         postprocess: Postprocessor | None = None,
     ) -> dict[Address, dict[str, Any]]:
         """Return typed predictions for a raw or encoded batch."""
-        supervised, _ = self.evaluate(
+
+        result = self.evaluate(
             batch=batch,
             preprocess=preprocess,
             postprocess=postprocess,
         )
-        return supervised
+
+        return result.predictions
 
     def embed(
         self,
@@ -560,12 +576,12 @@ class Model(lit.LightningModule):
         postprocess: Postprocessor | None = None,
     ) -> dict[Address, dict[str, Any]]:
         """Return configured embeddings for a raw or encoded batch."""
-        _, embeddings = self.evaluate(
+        result = self.evaluate(
             batch=batch,
             preprocess=preprocess,
             postprocess=postprocess,
         )
-        return embeddings
+        return result.embeddings
 
     training_step = partialmethod(step, strata=Strata.train)
     validation_step = partialmethod(step, strata=Strata.validate)

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/architecture/runtime.py

@@ -14,7 +14,8 @@ from tensordict import TensorDict
 from json2vec.architecture.contracts import sanitize
 from json2vec.architecture.encoder import ArrayEncoder
 from json2vec.architecture.node import NodeModule
-from json2vec.data.datasets.base import EncodedBatch
+from json2vec.data.datasets.base import EncodedBatch, EncodedInput
+from json2vec.data.iterables import encode
 from json2vec.structs.enums import Metric, Strata, TensorKey
 from json2vec.structs.packages import Embedding, Parcel, Prediction
 from json2vec.structs.tree import Address
@@ -50,9 +51,6 @@ class EvaluationResult:
     predictions: dict[Address, dict[str, Any]]
     embeddings: dict[Address, dict[str, Any]]
 
-    def as_tuple(self) -> tuple[dict[Address, dict[str, Any]], dict[Address, dict[str, Any]]]:
-        return self.predictions, self.embeddings
-
 
 class ModelRuntime:
     """Own runtime behavior that depends on an already-built model graph."""
@@ -182,16 +180,13 @@ class ModelRuntime:
         return supervised, embeddings
 
     @staticmethod
-    def evaluate(
+    def encode(
         module: "Model",
         batch: EncodedBatch | list[dict[str, Any]],
         preprocess: PreprocessFn | None = None,
-        postprocess: Postprocessor | None = None,
-    ) -> EvaluationResult:
-        from json2vec.data.iterables import encode
-
-        was_training = module.training
-        raw_batch = batch
+        strata: Strata | str = Strata.predict,
+    ) -> EncodedInput:
+        strata = Strata.normalize(strata)
 
         if preprocess is not None:
             observations: EncodedBatch = []
@@ -206,13 +201,24 @@ class ModelRuntime:
         elif batch and isinstance(batch[0], dict):
             batch = [[request] for request in cast(list[dict[str, Any]], batch)]
 
-        inputs = encode(
+        return encode(
             batch=cast(EncodedBatch, batch),
             hyperparameters=module.hyperparameters,
-            strata=Strata.predict,
+            strata=strata,
             interprocess_encoding_context=module.interprocess_encoding_context,
         )
 
+    @staticmethod
+    def evaluate(
+        module: "Model",
+        batch: EncodedBatch | list[dict[str, Any]],
+        preprocess: PreprocessFn | None = None,
+        postprocess: Postprocessor | None = None,
+    ) -> EvaluationResult:
+        was_training = module.training
+        raw_batch = batch
+        inputs = ModelRuntime.encode(module=module, batch=batch, preprocess=preprocess, strata=Strata.predict)
+
         module.eval()
         try:
             with torch.inference_mode():
@@ -226,7 +232,7 @@ class ModelRuntime:
         if postprocess is not None:
             context = {
                 "batch": raw_batch,
-                "observations": batch,
+                "observations": inputs[TensorKey.metadata],
                 "input": inputs,
                 TensorKey.metadata: inputs[TensorKey.metadata],
             }

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/structs/experiment.py

@@ -65,7 +65,7 @@ class Hyperparameters(Node):
     @classmethod
     def update_values(cls, values: Mapping[str, Any]) -> dict[str, Any]:
         normalized = dict(values)
-        target = normalized.pop("target", None)
+        target = normalized.get("target", None)
 
         if target is None:
             return normalized
@@ -76,11 +76,9 @@ class Hyperparameters(Node):
         if target:
             if normalized.get("p_prune") not in (None, 1.0):
                 raise ValueError("target=True is shorthand for p_prune=1.0")
-            normalized["p_prune"] = 1.0
         else:
-            if normalized.get("p_prune") is not None:
+            if "p_prune" in normalized and normalized["p_prune"] is not None:
                 raise ValueError("target=False is shorthand for p_prune=None")
-            normalized["p_prune"] = None
 
         return normalized
 
@@ -215,7 +213,7 @@ class Hyperparameters(Node):
     @property
     def target(self) -> list[Address]:
         role = NodePredicate(
-            func=lambda node: isinstance(node, Leaf) and node.active and getattr(node, "p_prune", 0.0) == 1.0,
+            func=lambda node: isinstance(node, Leaf) and node.active and node.target,
             key=("role", "target"),
         )
         return [Address(str(node.address)) for node in self.select(role)]
@@ -348,6 +346,8 @@ class Hyperparameters(Node):
 
             if validate and applicable_values:
                 payload = node.model_dump(mode="python", round_trip=True)
+                if "target" in applicable_values and "p_prune" not in applicable_values:
+                    payload.pop("p_prune", None)
                 payload.update(applicable_values)
                 validated = type(node).model_validate(payload)
                 applicable_values = {name: getattr(validated, name) for name in applicable_values}
@@ -501,7 +501,12 @@ class Hyperparameters(Node):
         nodes = self.select(*predicates, include_root=include_root, use_cache=use_cache)
         normalized_values = self.update_values(values)
         snapshot = [
-            (node, name, getattr(node, name, _MISSING), name in getattr(node, "model_fields_set", set()))
+            (
+                node,
+                "p_prune" if name == "target" else name,
+                getattr(node, "p_prune" if name == "target" else name, _MISSING),
+                ("p_prune" if name == "target" else name) in getattr(node, "model_fields_set", set()),
+            )
             for node in nodes
             for name in normalized_values
             if _has_model_attribute(node, name)

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/structs/selectors.py

@@ -140,7 +140,7 @@ class NodeAttribute(pydantic.BaseModel):
         if self.name == "descendants":
             return tuple(str(child.address) for child in getattr(node, "descendants", ()))
         if self.name == "target":
-            return isinstance(node, Leaf) and node.active and getattr(node, "p_prune", None) == 1.0
+            return isinstance(node, Leaf) and node.active and node.target
 
         extra = getattr(node, "model_extra", None) or {}
         if self.name in extra:

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/structs/tree.py

@@ -51,6 +51,17 @@ class Node(NodeMixin, pydantic.BaseModel):
     p_mask: Rate | None = None
     p_prune: PruneRate | None = None
 
+    @property
+    def target(self) -> bool:
+        return self.p_prune == 1.0
+
+    @target.setter
+    def target(self, value: bool) -> None:
+        if not isinstance(value, bool):
+            raise ValueError("target must be a boolean")
+
+        self.p_prune = 1.0 if value else None
+
     @classmethod
     def sanitize_name(cls, value: str) -> str:
         sanitized = re.sub(r"[^0-9A-Za-z_-]+", "_", value).strip("_")
@@ -78,6 +89,10 @@ class Node(NodeMixin, pydantic.BaseModel):
             if values.get("p_prune") not in (None, 1.0):
                 raise ValueError("target=True is shorthand for p_prune=1.0")
             values["p_prune"] = 1.0
+        else:
+            if values.get("p_prune") is not None:
+                raise ValueError("target=False is shorthand for p_prune=None")
+            values["p_prune"] = None
 
         return values
 

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/tensorfields/extensions/category.py

@@ -1,8 +1,9 @@
 # ty: ignore[invalid-method-override,unknown-argument]
 from __future__ import annotations
 
+from collections.abc import Mapping
 from functools import partial
-from typing import TYPE_CHECKING, Annotated, Literal, cast
+from typing import TYPE_CHECKING, Annotated, Any, Literal, cast
 
 import numpy as np
 import pydantic
@@ -46,10 +47,17 @@ class Request(RequestBase):
 
     type: Literal["category"] = "category"
     max_vocab_size: Annotated[int, pydantic.Field(gt=0, default=10_000)] = 10000
-    n_bands: Annotated[int, pydantic.Field(gt=0, default=8)] = 8
     p_unavailable: Annotated[float, pydantic.Field(ge=0.0, le=1.0, default=0.01)] = 0.01
     topk: list[int] | None = None
 
+    @pydantic.model_validator(mode="before")
+    @classmethod
+    def reject_removed_options(cls, data: Any) -> Any:
+        if isinstance(data, Mapping) and "n_bands" in data:
+            raise ValueError("Category does not support n_bands")
+
+        return data
+
     @pydantic.model_validator(mode="after")
     def check_topk(self):
         if self.topk is None:

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/tensorfields/extensions/dateparts.py

@@ -1,6 +1,7 @@
 # ty: ignore[invalid-argument-type,invalid-assignment,unknown-argument,unresolved-attribute]
 from __future__ import annotations
 
+import difflib
 import enum
 import math
 import re
@@ -74,6 +75,22 @@ class DatePart(enum.StrEnum):
         return cls.DEPTH[datepart]
 
 
+def _normalize_datepart_key(value: str) -> str:
+    value = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", value.strip())
+    value = re.sub(r"[^0-9A-Za-z]+", "_", value)
+    return value.strip("_").casefold()
+
+
+def _datepart_lookup() -> dict[str, DatePart]:
+    lookup: dict[str, DatePart] = {}
+    for datepart in DatePart:
+        normalized = _normalize_datepart_key(datepart.value)
+        lookup[normalized] = datepart
+        lookup[normalized.replace("_", "")] = datepart
+
+    return lookup
+
+
 @DatePart.day_of_month.register(depth=31)
 def _(arr: np.ndarray) -> np.ndarray:
     month_start = arr.astype("datetime64[M]")
@@ -133,6 +150,35 @@ class Request(RequestBase):
     dateparts: list[DatePart]
     pattern: Annotated[str | None, pydantic.Field(default=None)] = None
 
+    @pydantic.field_validator("dateparts", mode="before", check_fields=False)
+    @classmethod
+    def _coerce_dateparts(cls, value: Any) -> Any:
+        if not isinstance(value, (list, tuple)):
+            return value
+
+        lookup = _datepart_lookup()
+        canonical = [datepart.value for datepart in DatePart]
+        dateparts: list[DatePart] = []
+        for item in value:
+            if isinstance(item, DatePart):
+                dateparts.append(item)
+                continue
+
+            if not isinstance(item, str):
+                raise ValueError(f"datepart values must be strings, got {type(item).__name__}")
+
+            key = _normalize_datepart_key(item)
+            match = lookup.get(key) or lookup.get(key.replace("_", ""))
+            if match is not None:
+                dateparts.append(match)
+                continue
+
+            suggestions = difflib.get_close_matches(key, canonical, n=1)
+            suggestion = f"; did you mean '{suggestions[0]}'?" if suggestions else ""
+            raise ValueError(f"unknown datepart '{item}'{suggestion}")
+
+        return dateparts
+
     @pydantic.field_validator("dateparts", check_fields=False)
     @classmethod
     def check_dateparts(cls, v):

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/tensorfields/extensions/number.py

@@ -59,7 +59,7 @@ class Request(RequestBase):
     """Numeric scalar tensorfield request."""
 
     type: Literal["number"] = "number"
-    jitter: Annotated[float, pydantic.Field(ge=0.0, lt=1.0, default=0.0)] = 0.0
+    jitter: Annotated[float, pydantic.Field(ge=0.0, default=0.0)] = 0.0
     n_bands: Annotated[int, pydantic.Field(gt=0, default=8)] = 8
     offset: Annotated[int, pydantic.Field(gt=0, default=4)] = 4
     alpha: Annotated[float | None, pydantic.Field(gt=0.0, lt=1.0, default=None)] = None

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/tensorfields/extensions/set.py

@@ -43,6 +43,7 @@ class Request(RequestBase):
     type: Literal["set"] = "set"
     max_vocab_size: Annotated[int, pydantic.Field(gt=0, default=10_000)] = 10_000
     p_unavailable: Annotated[float, pydantic.Field(ge=0.0, le=1.0, default=0.01)] = 0.01
+    threshold: Annotated[float | None, pydantic.Field(ge=0.0, le=1.0, default=None)] = None
 
 
 def _items(value: Any) -> Iterable[Any]:
@@ -373,6 +374,7 @@ def loss(
 @sets.register
 def write(module: Model, prediction: Prediction):
     node = module.nodes[prediction.address]
+    request: Request = module.hyperparameters.requests[prediction.address]
     state_logits: torch.Tensor = prediction.payload[TensorKey.state]
     content_logits: torch.Tensor = prediction.payload[TensorKey.content]
 
@@ -383,7 +385,22 @@ def write(module: Model, prediction: Prediction):
 
     vocab = node.embedder.vocab.snapshot()
     probabilities = content_logits[..., : len(vocab)].sigmoid().detach().float().cpu().numpy()
-    content_payload = {str(label): probabilities[..., index] for index, label in enumerate(vocab)}
+    if request.threshold is None:
+        content_payload = {str(label): probabilities[..., index] for index, label in enumerate(vocab)}
+    else:
+        labels = np.asarray(vocab, dtype=object)
+
+        def pack_thresholded(values: np.ndarray) -> dict[str, float] | list:
+            if values.ndim == 1:
+                keep = values >= request.threshold
+                return {
+                    str(label): float(probability)
+                    for label, probability in zip(labels[keep].tolist(), values[keep].tolist())
+                }
+
+            return [pack_thresholded(values[index]) for index in range(values.shape[0])]
+
+        content_payload = pack_thresholded(probabilities)
 
     return {
         TensorKey.state.name: state_payload,

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec/tensorfields/extensions/vector.py

@@ -226,7 +226,11 @@ class Decoder(DecoderBase):
 
         request: Request = hyperparameters.requests[address]
 
-        self.linear = torch.nn.Linear(
+        self.classification = torch.nn.Linear(
+            in_features=hyperparameters.d_model,
+            out_features=len(Tokens),
+        )
+        self.regression = torch.nn.Linear(
             in_features=hyperparameters.d_model,
             out_features=request.n_dim,
         )
@@ -235,7 +239,8 @@ class Decoder(DecoderBase):
     def decode(self, pooled: torch.Tensor) -> TensorDict[TensorKey, torch.Tensor]:
         return TensorDict(
             source={
-                TensorKey.content: self.linear(pooled),
+                TensorKey.state: self.classification(pooled),
+                TensorKey.content: self.regression(pooled),
             }
         )
 
@@ -251,30 +256,67 @@ def loss(
     request: Request = module.hyperparameters.requests[address]
 
     trainable = batch.trainable.reshape(-1)
+    state_targets = batch.targets[TensorKey.state].reshape(-1)
+    state_inputs = prediction.payload[TensorKey.state].reshape(-1, len(Tokens))
+
+    output: torch.Tensor = module.track(
+        (address, strata, Metric.loss, TensorKey.state),
+        value=(
+            torch.nn.functional.cross_entropy(
+                input=state_inputs,
+                target=state_targets,
+                reduction="none",
+            )
+            .masked_select(trainable)
+            .mean()
+        ),
+    )
+
+    module.track(
+        (address, strata, Metric.accuracy, TensorKey.state),
+        value=state_inputs.argmax(dim=1).eq(state_targets).masked_select(trainable).float().mean(),
+    )
+
+    valued = trainable & state_targets.eq(Tokens.valued.value)
+    if not valued.any():
+        return output
+
     inputs = prediction.payload[TensorKey.content].reshape(-1, request.n_dim)
     targets = batch.targets[TensorKey.content].reshape(-1, request.n_dim)
     diff = inputs.subtract(targets)
 
-    loss: torch.Tensor = module.track(
+    output += module.track(
         (address, strata, Metric.loss, TensorKey.content),
-        value=request.objective.loss(inputs=inputs, targets=targets).masked_select(trainable).mean(),
+        value=request.objective.loss(inputs=inputs, targets=targets).masked_select(valued).mean(),
     )
 
     module.track(
         (address, strata, Metric.mae, TensorKey.content),
-        value=diff.absolute().mean(dim=1).masked_select(trainable).mean(),
+        value=diff.absolute().mean(dim=1).masked_select(valued).mean(),
     )
 
     module.track(
         (address, strata, Metric.rmse, TensorKey.content),
-        value=diff.square().mean(dim=1).sqrt().masked_select(trainable).mean(),
+        value=diff.square().mean(dim=1).sqrt().masked_select(valued).mean(),
     )
 
-    return loss
+    return output
 
 
 @vector.register
 def write(module: Model, prediction: Prediction):
+    content: np.ndarray = prediction.payload[TensorKey.content].detach().float().cpu().numpy()
+    state_logits: torch.Tensor = prediction.payload[TensorKey.state]
+    tokens: np.ndarray = np.fromiter((token.name for token in Tokens), dtype=object, count=len(Tokens))
+    state_log_norm = state_logits.logsumexp(dim=-1, keepdim=True)
+    state_distribution = (state_logits - state_log_norm).exp().detach().float().cpu().numpy()
+    state_payload = {token: state_distribution[..., index] for index, token in enumerate(tokens.tolist())}
+
+    non_valued = state_logits.argmax(dim=-1).ne(Tokens.valued.value).detach().cpu().numpy()
+    content = content.copy()
+    content[non_valued] = 0.0
+
     return {
-        TensorKey.content.name: prediction.payload[TensorKey.content].detach().float().cpu().numpy(),
+        TensorKey.state.name: state_payload,
+        TensorKey.content.name: content,
     }

{json2vec-0.4.4 → json2vec-0.4.5/src/json2vec.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.4
+Version: 0.4.5
 Summary: {...} -> [*]
 License-Expression: Apache-2.0
 Requires-Python: >=3.12

{json2vec-0.4.4 → json2vec-0.4.5}/src/json2vec.egg-info/SOURCES.txt

@@ -15,13 +15,13 @@ src/json2vec/architecture/checkpoint.py
 src/json2vec/architecture/contracts.py
 src/json2vec/architecture/encoder.py
 src/json2vec/architecture/graph.py
+src/json2vec/architecture/mutations.py
 src/json2vec/architecture/node.py
 src/json2vec/architecture/plot.py
 src/json2vec/architecture/pool.py
 src/json2vec/architecture/root.py
 src/json2vec/architecture/rotary.py
 src/json2vec/architecture/runtime.py
-src/json2vec/architecture/schema_editor.py
 src/json2vec/data/__init__.py
 src/json2vec/data/iterables.py
 src/json2vec/data/processing.py

json2vec-0.4.4/src/json2vec/architecture/schema_editor.py

@@ -1,126 +0,0 @@
-"""Model-facing schema mutation orchestration."""
-
-from __future__ import annotations
-
-from collections.abc import Callable, Iterator
-from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any
-
-from json2vec.architecture.graph import ModelGraph
-from json2vec.structs.experiment import NodeAttribute, NodePredicate, SchemaField
-from json2vec.structs.tree import Node
-
-if TYPE_CHECKING:
-    from json2vec.architecture.root import Model
-
-
-class SchemaEditor:
-    """Coordinate schema mutations with runtime graph rebuilds."""
-
-    def __init__(self, module: "Model") -> None:
-        self.module = module
-
-    def select(
-        self,
-        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
-        include_root: bool = True,
-        use_cache: bool = True,
-    ) -> list[Node]:
-        return self.module.hyperparameters.select(
-            *predicates,
-            include_root=include_root,
-            use_cache=use_cache,
-        )
-
-    def update(
-        self,
-        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
-        strict: bool = True,
-        allow_extra: bool = False,
-        include_root: bool = True,
-        validate: bool = True,
-        use_cache: bool = False,
-        **values: Any,
-    ) -> None:
-        self.module._assert_mutation_allowed("update")
-        self.module.hyperparameters.update(
-            *predicates,
-            strict=strict,
-            allow_extra=allow_extra,
-            include_root=include_root,
-            validate=validate,
-            use_cache=use_cache,
-            **values,
-        )
-        ModelGraph.rebuild(self.module)
-        self.module._reset_contracts()
-
-    def extend(
-        self,
-        *args: NodePredicate | NodeAttribute | Callable[[Node], bool] | SchemaField,
-        include_root: bool = True,
-        use_cache: bool = True,
-    ) -> None:
-        self.module._assert_mutation_allowed("extend")
-        self.module.hyperparameters.extend(*args, include_root=include_root, use_cache=use_cache)
-        ModelGraph.rebuild(self.module)
-        self.module._reset_contracts()
-
-    def delete(
-        self,
-        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
-        include_root: bool = False,
-        use_cache: bool = True,
-    ) -> None:
-        self.module._assert_mutation_allowed("delete")
-        self.module.hyperparameters.delete(*predicates, include_root=include_root, use_cache=use_cache)
-        ModelGraph.rebuild(self.module)
-        self.module._reset_contracts()
-
-    def reset(
-        self,
-        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
-        include_root: bool = True,
-        use_cache: bool = True,
-        descendants: bool = False,
-    ) -> None:
-        self.module._assert_mutation_allowed("reset")
-        selected = self.module.hyperparameters.select(
-            *predicates,
-            include_root=include_root,
-            use_cache=use_cache,
-        )
-        if not selected:
-            raise ValueError("reset matched no nodes")
-
-        ModelGraph.reset_selected(self.module, selected, descendants=descendants)
-        self.module._reset_contracts()
-
-    @contextmanager
-    def override(
-        self,
-        *predicates: NodePredicate | NodeAttribute | Callable[[Node], bool],
-        strict: bool = True,
-        allow_extra: bool = False,
-        include_root: bool = True,
-        validate: bool = True,
-        use_cache: bool = False,
-        **values: Any,
-    ) -> Iterator[None]:
-        self.module._assert_mutation_allowed("override")
-        try:
-            with self.module.hyperparameters.override(
-                *predicates,
-                strict=strict,
-                allow_extra=allow_extra,
-                include_root=include_root,
-                validate=validate,
-                use_cache=use_cache,
-                **values,
-            ):
-                ModelGraph.rebuild(self.module)
-                self.module._reset_contracts()
-                yield
-        finally:
-            ModelGraph.rebuild(self.module)
-            self.module._reset_contracts()