json2vec 0.4.0__tar.gz → 0.4.3__tar.gz

{json2vec-0.4.0/src/json2vec.egg-info → json2vec-0.4.3}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: json2vec
-Version: 0.4.0
-Summary: JSON -> [*]
+Version: 0.4.3
+Summary: {...} -> [*]
 License-Expression: Apache-2.0
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown

{json2vec-0.4.0 → json2vec-0.4.3}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "json2vec"
-version = "0.4.0"
-description = "JSON -> [*]"
+version = "0.4.3"
+description = "{...} -> [*]"
 readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.12"

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/architecture/encoder.py

@@ -13,6 +13,7 @@ from json2vec.structs.tree import Address
 if TYPE_CHECKING:
     from json2vec.structs.experiment import Hyperparameters
 
+
 class RotaryTransformerEncoderLayer(torch.nn.Module):
     def __init__(
         self,

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/architecture/node.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import torch
 
@@ -25,7 +25,7 @@ class NodeModule(torch.nn.Module):
         if address in hyperparameters.requests:
             request: Node = hyperparameters.requests[address]
             plugin: Plugin = TENSORFIELDS[request.type]
-            embedder_kwargs = dict(hyperparameters=hyperparameters, address=address)
+            embedder_kwargs: dict[str, Any] = dict(hyperparameters=hyperparameters, address=address)
             if "batch_size" in inspect.signature(plugin.Embedder.__init__).parameters:
                 embedder_kwargs["batch_size"] = batch_size
 

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/architecture/plot.py

@@ -49,7 +49,7 @@ def plot(
 ) -> None:
     """Print a Rich model visualization and optionally write it as text."""
     renderable = build_plot(module=module, address=address, detail=detail, mode=mode)
-    Console(width=PLOT_WIDTH, force_jupyter=False).print(renderable)
+    Console(width=PLOT_WIDTH).print(renderable)
 
     if out is None:
         return
@@ -112,7 +112,9 @@ def render_flow_plot(module: "Model", address: Address | str | None) -> Renderab
     table.add_column("Count", justify="right")
     table.add_row("JSON", "Raw nested records enter with the shape described by the schema.", "")
     table.add_row("Tensorfields", "Typed requests read values with JMESPath queries.", str(len(fields)))
-    table.add_row("Encoders", "Array nodes pool child embeddings into parent contexts.", str(len(hyperparameters.arrays)))
+    table.add_row(
+        "Encoders", "Array nodes pool child embeddings into parent contexts.", str(len(hyperparameters.arrays))
+    )
     table.add_row("Targets", "Target fields produce supervised predictions.", str(target_count))
     table.add_row("Embeddings", "Selected nodes expose reusable embeddings.", str(embed_count))
 
@@ -272,11 +274,7 @@ def node_metadata_keys(node: Node, values: dict[str, Any], state_focus: bool) ->
 
 
 def should_hide_metadata(key: str, value: Any) -> bool:
-    return (
-        (key == "active" and value is True)
-        or (key == "embed" and value is False)
-        or key == "description"
-    )
+    return (key == "active" and value is True) or (key == "embed" and value is False) or key == "description"
 
 
 def format_metadata_value(value: Any) -> str:
@@ -317,7 +315,9 @@ def format_detail_inline(value: Any) -> str:
         return truncate(value, width=100)
 
     if isinstance(value, list):
-        return truncate(format_inline_sequence(value) or pformat(value, compact=True, sort_dicts=False, width=88), width=100)
+        return truncate(
+            format_inline_sequence(value) or pformat(value, compact=True, sort_dicts=False, width=88), width=100
+        )
 
     return truncate(pformat(value, compact=True, sort_dicts=False, width=88), width=100)
 
@@ -332,7 +332,9 @@ def summarize_value(value: Any, max_items: int = 8) -> Any:
     if isinstance(value, list):
         if len(value) <= max_items:
             return [summarize_value(item, max_items=max_items) for item in value]
-        return [summarize_value(item, max_items=max_items) for item in value[:max_items]] + [f"... {len(value) - max_items} more"]
+        return [summarize_value(item, max_items=max_items) for item in value[:max_items]] + [
+            f"... {len(value) - max_items} more"
+        ]
 
     return value
 
@@ -410,8 +412,9 @@ def format_compact_number(value: Any) -> str:
 
 def resolve_node(hyperparameters: "Hyperparameters", address: Address | str) -> Node:
     key = Address(str(address))
-    leaves = {node.address: node for node in hyperparameters.descendants if isinstance(node, Leaf)}
-    nodes: dict[Address, Node] = hyperparameters.arrays | leaves
+    leaves: dict[Address, Node] = {node.address: node for node in hyperparameters.descendants if isinstance(node, Leaf)}
+    nodes: dict[Address, Node] = dict(hyperparameters.arrays)
+    nodes.update(leaves)
 
     if key not in nodes:
         raise ValueError(f"address '{address}' was not found in the hyperparameters")

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/architecture/root.py

@@ -18,6 +18,7 @@ from tensordict import TensorDict
 
 from json2vec.architecture.encoder import ArrayEncoder
 from json2vec.architecture.node import NodeModule
+from json2vec.architecture.plot import PlotMode
 from json2vec.data.datasets.base import EncodedBatch
 from json2vec.data.iterables import encode, mock
 from json2vec.structs.enums import AttentionMode, Metric, Strata, TensorKey
@@ -296,7 +297,6 @@ class Model(lit.LightningModule):
         optimizer: OptimizerConfig | None = None,
         scheduler: SchedulerConfig | None = None,
     ):
-
         super().__init__()
         if batch_size <= 0:
             raise ValueError("batch_size must be > 0")
@@ -330,11 +330,16 @@ class Model(lit.LightningModule):
         self.example_input_array = mock(hyperparameters=self.hyperparameters, batch_size=self.batch_size)
 
     def _rebuild(self) -> None:
+        self.hyperparameters._clear_tree_caches()
+        was_training = self.training
+        device = self.device
         previous = {
             name: value.detach().clone() if isinstance(value, torch.Tensor) else deepcopy(value)
             for name, value in self.state_dict().items()
         }
         self._build()
+        if isinstance(device, torch.device):
+            self.to(device=device)
         current = self.state_dict()
         compatible = {}
         for name, value in previous.items():
@@ -351,6 +356,7 @@ class Model(lit.LightningModule):
             compatible[name] = value
 
         self.load_state_dict(compatible, strict=False)
+        self.train(was_training)
 
     def select(
         self,
@@ -372,6 +378,7 @@ class Model(lit.LightningModule):
         allow_extra: bool = False,
         include_root: bool = True,
         validate: bool = True,
+        use_cache: bool = False,
         **values: Any,
     ) -> None:
         """Mutate selected schema nodes and rebuild compatible modules.
@@ -386,6 +393,8 @@ class Model(lit.LightningModule):
                 allow unknown fields.
             include_root: Include the root node in predicate matching.
             validate: Validate each node after applying candidate values.
+            use_cache: Permit cached selector results. Mutations default this to
+                `False` so updates always evaluate against current schema state.
             **values: Schema attributes to update.
         """
         self._assert_mutation_allowed("update")
@@ -395,6 +404,7 @@ class Model(lit.LightningModule):
             allow_extra=allow_extra,
             include_root=include_root,
             validate=validate,
+            use_cache=use_cache,
             **values,
         )
         self._rebuild()
@@ -468,6 +478,7 @@ class Model(lit.LightningModule):
         allow_extra: bool = False,
         include_root: bool = True,
         validate: bool = True,
+        use_cache: bool = False,
         **values: Any,
     ) -> Iterator[None]:
         """Temporarily mutate selected schema nodes and keep runtime modules synchronized."""
@@ -479,6 +490,7 @@ class Model(lit.LightningModule):
                 allow_extra=allow_extra,
                 include_root=include_root,
                 validate=validate,
+                use_cache=use_cache,
                 **values,
             ):
                 self._rebuild()
@@ -563,7 +575,7 @@ class Model(lit.LightningModule):
         address: Address | str | None = None,
         detail: bool = False,
         out: str | Path | None = None,
-        mode: str = "schema",
+        mode: PlotMode = "schema",
     ) -> None:
         """Print a Rich model visualization.
 
@@ -602,7 +614,8 @@ class Model(lit.LightningModule):
 
             embedder: EmbedderBase = self.nodes[address].embedder
             embedding: Parcel = embedder(tensorfield)
-            processed[embedding.destination].append(embedding)
+            if embedding.destination is not None:
+                processed[embedding.destination].append(embedding)
             outgoing[embedding.origin] = embedding
 
             if address in self.hyperparameters.embed:
@@ -617,7 +630,8 @@ class Model(lit.LightningModule):
 
                 encoder: ArrayEncoder = self.nodes[address].encoder
                 encoding: Parcel = encoder(processed[address])
-                processed[encoding.destination].append(encoding)
+                if encoding.destination is not None:
+                    processed[encoding.destination].append(encoding)
                 outgoing[encoding.origin] = encoding
 
                 if address in self.hyperparameters.embed:
@@ -696,7 +710,6 @@ class Model(lit.LightningModule):
     def write(
         self, predictions: list[Prediction]
     ) -> tuple[dict[Address, dict[str, Any]], dict[Address, dict[str, Any]]]:
-
         supervised: dict[Address, dict[str, Any]] = {}
         embeddings: dict[Address, dict[str, Any]] = {}
 
@@ -738,17 +751,21 @@ class Model(lit.LightningModule):
 
         if preprocess is not None:
             observations: EncodedBatch = []
-            for request in batch:
+            for request in cast(list[dict[str, Any]], batch):
                 observation = preprocess(request)
                 if not isinstance(observation, dict):
                     raise TypeError(f"preprocessor must return a dict object, got {type(observation).__name__}")
 
                 observations.append([observation])
 
-            batch = observations
+            encoded_batch = observations
+        elif batch and isinstance(batch[0], dict):
+            encoded_batch = [[request] for request in cast(list[dict[str, Any]], batch)]
+        else:
+            encoded_batch = cast(EncodedBatch, batch)
 
         inputs = encode(
-            batch=batch,
+            batch=encoded_batch,
             hyperparameters=self.hyperparameters,
             strata=Strata.predict,
             interprocess_encoding_context=self.interprocess_encoding_context,

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/architecture/rotary.py

@@ -13,6 +13,7 @@ class RotaryEmbedding(torch.nn.Module):
         self.base = base
 
         index = torch.arange(0, self.rotary_dim, 2, dtype=torch.float32)
+        self.inv_freq: torch.Tensor
         self.register_buffer("inv_freq", base ** (-index / self.rotary_dim), persistent=False)
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/data/datasets/polars.py

@@ -57,7 +57,7 @@ def _dataframes_by_strata(dataframe: pl.DataFrame | DataFrameMap) -> dict[Strata
         return {strata: dataframe for strata in Strata}
 
     normalized: dict[Strata, pl.DataFrame] = {}
-    for key, frame in dataframe.items():
+    for key, frame in cast(DataFrameMap, dataframe).items():
         if not isinstance(frame, pl.DataFrame):
             raise TypeError(f"dataframe for strata '{key}' must be a polars DataFrame")
         normalized[Strata.normalize(key)] = frame
@@ -79,7 +79,7 @@ def observe_polars(
     world_size: int | None = None,
 ) -> Iterator[RawObservation]:
     if replacement:
-        rows = cast(list[RawObservation], dataframe.to_dicts())
+        rows = dataframe.to_dicts()
         if not rows:
             raise ValueError("no dataframe rows available for replacement sampling")
 
@@ -103,7 +103,7 @@ def observe_polars(
                 worker_id=worker_id,
                 num_workers=num_workers,
             ):
-                yield cast(RawObservation, row)
+                yield row
         return
 
     for chunk_index, offset in enumerate(range(0, dataframe.height, chunk_batch_size)):
@@ -115,7 +115,7 @@ def observe_polars(
             ):
                 continue
 
-        yield from cast(list[RawObservation], dataframe.slice(offset, chunk_batch_size).to_dicts())
+        yield from dataframe.slice(offset, chunk_batch_size).to_dicts()
 
 
 class PolarsBatchDataset(IterableDataset):
@@ -281,7 +281,6 @@ class PolarsDataModule(lit.LightningDataModule):
         else:
             dataframes = _dataframes_by_strata(dataframe)
 
-        self.hyperparameters = model.hyperparameters
         self.dataframes = dataframes
         self.preprocessor = PreprocessorConfig.normalize(preprocessor)
         self.preprocessor_kwargs = dict(kwargs)
@@ -289,26 +288,55 @@ class PolarsDataModule(lit.LightningDataModule):
             self._model_ref = weakref.ref(model)
         except TypeError:
             self._model_ref = None
+        self._hyperparameters = model.hyperparameters
         self._interprocess_encoding_context = model.interprocess_encoding_context
-        self.batch_size = model.batch_size
+        self._batch_size = model.batch_size
         self.num_workers = Strata.expand(num_workers, default=None)
         self.persistent_workers = Strata.expand(persistent_workers, default=True)
         self.pin_memory = Strata.expand(pin_memory, default=True)
         self.sharding = ShardingStrategy.expand(sharding, default=ShardingStrategy.chunk)
         self.chunk_batch_size = Strata.expand(chunk_batch_size, default=4096)
         self.observation_buffer_size = Strata.expand(observation_buffer_size, default=1)
-        self.sample_rate = {
-            strata: float(rate)
-            for strata, rate in Strata.expand(sample_rate, default=1.0).items()
-        }
+        self.sample_rate = {strata: float(rate) for strata, rate in Strata.expand(sample_rate, default=1.0).items()}
         self.replacement = Strata.expand(replacement, default=False)
 
+    def _model(self) -> Model | None:
+        if self._model_ref is None:
+            return None
+
+        return self._model_ref()
+
+    @property
+    def hyperparameters(self) -> Hyperparameters:
+        model = self._model()
+        if model is not None:
+            return model.hyperparameters
+
+        return self._hyperparameters
+
+    @hyperparameters.setter
+    def hyperparameters(self, hyperparameters: Hyperparameters) -> None:
+        self._model_ref = None
+        self._hyperparameters = hyperparameters
+
+    @property
+    def batch_size(self) -> int:
+        model = self._model()
+        if model is not None:
+            return model.batch_size
+
+        return self._batch_size
+
+    @batch_size.setter
+    def batch_size(self, batch_size: int) -> None:
+        self._model_ref = None
+        self._batch_size = batch_size
+
     @property
     def interprocess_encoding_context(self) -> InterprocessEncodingContext:
-        if self._model_ref is not None:
-            model = self._model_ref()
-            if model is not None:
-                return model.interprocess_encoding_context
+        model = self._model()
+        if model is not None:
+            return model.interprocess_encoding_context
 
         return self._interprocess_encoding_context
 

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/data/datasets/streaming.py

@@ -118,8 +118,7 @@ def observe(
         sampled_paths = list(paths)
         if not sampled_paths:
             raise ValueError(
-                "no matching files available for replacement sampling; "
-                "check the streaming root and split pattern"
+                "no matching files available for replacement sampling; check the streaming root and split pattern"
             )
 
         def choices() -> Iterator[str]:
@@ -405,7 +404,6 @@ class StreamingDataModule(lit.LightningDataModule):
     ):
         super().__init__()
 
-        self.hyperparameters = model.hyperparameters
         self.root = root
         self.suffix = Suffix(suffix)
         self.train = train
@@ -418,8 +416,9 @@ class StreamingDataModule(lit.LightningDataModule):
             self._model_ref = weakref.ref(model)
         except TypeError:
             self._model_ref = None
+        self._hyperparameters = model.hyperparameters
         self._interprocess_encoding_context = model.interprocess_encoding_context
-        self.batch_size = model.batch_size
+        self._batch_size = model.batch_size
         self.num_workers = Strata.expand(num_workers, default=None)
         self.persistent_workers = Strata.expand(persistent_workers, default=True)
         self.pin_memory = Strata.expand(pin_memory, default=True)
@@ -427,22 +426,50 @@ class StreamingDataModule(lit.LightningDataModule):
         self.chunk_batch_size = Strata.expand(chunk_batch_size, default=4096)
         self.file_buffer_size = Strata.expand(file_buffer_size, default=1)
         self.observation_buffer_size = Strata.expand(observation_buffer_size, default=1)
-        self.sample_rate = {
-            strata: float(rate)
-            for strata, rate in Strata.expand(sample_rate, default=1.0).items()
-        }
+        self.sample_rate = {strata: float(rate) for strata, rate in Strata.expand(sample_rate, default=1.0).items()}
         self.replacement = (
             {strata: strata == Strata.train for strata in Strata}
             if replacement is None
             else Strata.expand(replacement, default=False)
         )
 
+    def _model(self) -> Model | None:
+        if self._model_ref is None:
+            return None
+
+        return self._model_ref()
+
+    @property
+    def hyperparameters(self) -> Hyperparameters:
+        model = self._model()
+        if model is not None:
+            return model.hyperparameters
+
+        return self._hyperparameters
+
+    @hyperparameters.setter
+    def hyperparameters(self, hyperparameters: Hyperparameters) -> None:
+        self._model_ref = None
+        self._hyperparameters = hyperparameters
+
+    @property
+    def batch_size(self) -> int:
+        model = self._model()
+        if model is not None:
+            return model.batch_size
+
+        return self._batch_size
+
+    @batch_size.setter
+    def batch_size(self, batch_size: int) -> None:
+        self._model_ref = None
+        self._batch_size = batch_size
+
     @property
     def interprocess_encoding_context(self) -> InterprocessEncodingContext:
-        if self._model_ref is not None:
-            model = self._model_ref()
-            if model is not None:
-                return model.interprocess_encoding_context
+        model = self._model()
+        if model is not None:
+            return model.interprocess_encoding_context
 
         return self._interprocess_encoding_context
 

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/data/processing.py

@@ -141,7 +141,7 @@ class Pipeline:
         return self
 
     def __repr__(self):
-        return f"Pipeline({repr(self.source)}, {repr(self.arguments)})"
+        return f"Pipeline(steps={len(self.steps)}, arguments={self.arguments!r})"
 
     def __iter__(self):
         stream = self.steps[0]()

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/inference/callback.py

@@ -41,9 +41,7 @@ class Writer(callbacks.BasePredictionWriter):
         self.writer: pq.ParquetWriter | None = None
 
     @staticmethod
-    def _as_struct_frame(
-        values_by_address: dict[Address, dict[str, Any]], alias: str, num_rows: int
-    ) -> pl.DataFrame:
+    def _as_struct_frame(values_by_address: dict[Address, dict[str, Any]], alias: str, num_rows: int) -> pl.DataFrame:
         if len(values_by_address) == 0:
             return pl.DataFrame({alias: [None] * num_rows})
 
@@ -64,7 +62,7 @@ class Writer(callbacks.BasePredictionWriter):
         batch: TensorDict[Address, TensorFieldBase],
         batch_idx: int,
         dataloader_idx: int,
-    ) -> None:
+    ) -> None:  # ty:ignore[invalid-method-override]
         num_rows = len(batch["metadata"])
 
         supervised: dict[Address, dict[str, Any]]

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/inference/deployment.py

@@ -4,7 +4,7 @@ import functools
 from collections.abc import Callable
 from enum import StrEnum
 from pathlib import Path
-from typing import Any, TypeAlias
+from typing import Any, TypeAlias, cast
 
 import litserve as ls
 import pydantic
@@ -42,7 +42,7 @@ class Accelerator(StrEnum):
         if normalized == "":
             raise ValueError("accelerator must not be blank")
 
-        return cls._value2member_map_.get(normalized)
+        return cast(Accelerator | None, cls._value2member_map_.get(normalized))
 
 
 class ErrorItem(pydantic.BaseModel):
@@ -114,8 +114,7 @@ class API(ls.LitAPI):
         self,
         request: dict[str, Any] | pydantic.BaseModel,
         context: dict[str, Any] | None = None,
-    ) -> Input | ErrorItem:
-
+    ) -> Input | ErrorItem:  # ty:ignore[invalid-method-override]
         if isinstance(request, pydantic.BaseModel):
             request = request.model_dump()
 
@@ -175,13 +174,13 @@ class API(ls.LitAPI):
         return BatchItem(data=data, valid_indices=valid_indices, items=inputs)
 
     @beartype
-    def unbatch(self, outputs: list[Any]) -> list[Any]:
+    def unbatch(self, outputs: list[Any]) -> list[Any]:  # ty:ignore[invalid-method-override]
         return list(outputs)
 
     @beartype
     def predict(
         self, data: BatchItem | Input | ErrorItem
-    ) -> list[list[Prediction] | ErrorItem] | list[Prediction] | ErrorItem:
+    ) -> list[list[Prediction] | ErrorItem] | list[Prediction] | ErrorItem:  # ty:ignore[invalid-method-override]
         if isinstance(data, ErrorItem):
             return data
 
@@ -209,7 +208,7 @@ class API(ls.LitAPI):
         self,
         response: list[Prediction] | ErrorItem,
         context: dict[str, Any] | None = None,
-    ) -> dict[str, Any] | pydantic.BaseModel:
+    ) -> dict[str, Any] | pydantic.BaseModel:  # ty:ignore[invalid-method-override]
         if isinstance(response, ErrorItem):
             return {
                 "predictions": {},

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/logging/config.py

@@ -10,6 +10,7 @@ console = Console(file=sys.stdout)
 
 LOG_LEVEL: str = os.getenv("JSON2VEC_LOG_LEVEL", "DEBUG").upper()
 
+
 def sink(message):
     record = message.record
     extras = {k: str(v) for k, v in record["extra"].items()}

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/preprocessors/base.py

@@ -50,8 +50,7 @@ class Preprocessor(pydantic.BaseModel):
     def accepted_kwargs(func: Callable[..., Any]) -> tuple[bool, frozenset[str]]:
         signature = inspect.signature(func)
         accepts_variadic_kwargs = any(
-            parameter.kind == inspect.Parameter.VAR_KEYWORD
-            for parameter in signature.parameters.values()
+            parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()
         )
         accepted = frozenset(signature.parameters.keys())
         return accepts_variadic_kwargs, accepted
@@ -66,7 +65,8 @@ class Preprocessor(pydantic.BaseModel):
 
     @classmethod
     def register(cls, func: Callable[..., Any], *, mode: PreprocessorMode) -> Callable[..., Any]:
-        PREPROCESSORS[func.__name__] = cls(name=func.__name__, func=func, mode=mode)
+        name = getattr(func, "__name__", type(func).__name__)
+        PREPROCESSORS[name] = cls(name=name, func=func, mode=mode)
         return func
 
     def __call__(self, observation: dict, **kwargs) -> Any:
@@ -99,9 +99,7 @@ class Preprocessor(pydantic.BaseModel):
 
     def require_object(self, output: Any, *, mode: PreprocessorMode) -> dict[str, Any]:
         if not isinstance(output, dict):
-            raise TypeError(
-                f"{mode} preprocessor '{self.name}' must produce dict objects, got {type(output).__name__}"
-            )
+            raise TypeError(f"{mode} preprocessor '{self.name}' must produce dict objects, got {type(output).__name__}")
 
         return output
 

{json2vec-0.4.0 → json2vec-0.4.3}/src/json2vec/structs/enums.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 import enum
 from collections.abc import Mapping
-from typing import TypeVar
+from typing import TypeVar, cast
 
 T = TypeVar("T")
+DefaultT = TypeVar("DefaultT")
 
 
 class Tokens(enum.IntEnum):
@@ -27,10 +30,10 @@ class Strata(enum.StrEnum):
         return cls(str(value).strip().lower())
 
     @classmethod
-    def expand(cls, value: T | Mapping["Strata | str", T], *, default: T) -> dict["Strata", T]:
+    def expand(cls, value: T | Mapping[Strata | str, T], *, default: DefaultT) -> dict[Strata, T | DefaultT]:
         if isinstance(value, Mapping):
-            normalized = {strata: default for strata in cls}
-            for key, item in value.items():
+            normalized: dict[Strata, T | DefaultT] = {strata: default for strata in cls}
+            for key, item in cast(Mapping[Strata | str, T], value).items():
                 normalized[cls.normalize(key)] = item
             return normalized
 
@@ -87,10 +90,7 @@ class ShardingStrategy(enum.StrEnum):
         *,
         default: "ShardingStrategy",
     ) -> dict[Strata, "ShardingStrategy"]:
-        return {
-            strata: cls.normalize(strategy)
-            for strata, strategy in Strata.expand(value, default=default).items()
-        }
+        return {strata: cls.normalize(strategy) for strata, strategy in Strata.expand(value, default=default).items()}
 
 
 class AttentionMode(enum.StrEnum):