json2vec 0.1.0__py3-none-any.whl

json2vec/tensorfields/base.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+import functools
+import math
+import re
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Type, TypeAlias
+
+import pluggy
+import torch
+from tensordict import TensorDict
+
+from json2vec.architecture.pool import LearnedQueryCrossAttention
+from json2vec.structs.enums import Component, Metric, Strata, TensorKey
+from json2vec.structs.packages import Parcel, Prediction
+from json2vec.structs.tree import Address, Leaf, Node
+from json2vec.tensorfields.spec import PluginSpec
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import JSON2Vec
+    from json2vec.structs.experiment import Session
+    from json2vec.structs.structure import Structure
+
+pm: pluggy.PluginManager = pluggy.PluginManager(project_name="tensorfields")
+
+pm.add_hookspecs(module_or_class=PluginSpec)
+
+RequestBase: TypeAlias = Leaf
+
+
+class EmbedderBase(torch.nn.Module):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__()
+
+
+class DecoderBase(torch.nn.Module):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__()
+
+        self.address: Address = address
+        self.sigma: torch.Tensor = torch.nn.Parameter(torch.zeros(1))
+
+        self.pool = LearnedQueryCrossAttention(
+            n_context=math.prod(structure.shapes[address]),
+            d_model=structure.d_model,
+            nhead=structure.requests[address].n_heads,
+            dropout=structure.dropout,
+            n_linear=structure.requests[address].n_linear,
+        )
+
+    def decode(self, pooled: torch.Tensor) -> TensorDict[TensorKey, torch.Tensor]:
+        raise NotImplementedError("decoder must implement decode(pooled)")
+
+    def forward(self, parcels: list[Parcel]) -> Prediction:
+        if len(parcels) == 0:
+            raise ValueError("decoder requires at least one parcel")
+
+        N, *_, C = parcels[0].payload.shape
+        stacked = torch.cat([parcel.payload.reshape(N, -1, C) for parcel in parcels], dim=1)
+        pooled = self.pool(stacked)
+
+        payload = self.decode(pooled)
+        return Prediction(
+            payload=payload,
+            address=self.address,
+            batch_size=pooled.shape[0],
+        )
+
+
+class TensorFieldBase(ABC):
+    content: torch.Tensor
+    state: torch.Tensor
+    trainable: torch.Tensor
+    targets: TensorDict[TensorKey, torch.Tensor]
+
+    @classmethod
+    @abstractmethod
+    def new(
+        cls,
+        values: list,
+        address: Address,
+        session: Session,
+        strata: Strata,
+        state: Any,
+    ) -> "TensorFieldBase":
+        raise NotImplementedError
+
+    @abstractmethod
+    def mask(self, p_mask: float):
+        raise NotImplementedError
+
+    @abstractmethod
+    def prune(cls, p_prune: float):
+        raise NotImplementedError
+
+
+TENSORFIELDS: dict[str, "Plugin"] = {}
+
+
+class Plugin:
+    def __init__(self, name: str):
+        if not isinstance(name, str):
+            raise TypeError("Plugin name must be a string")
+
+        # should start with a letter and contain only lowercase letters, numbers, and underscores
+        if not re.match(r"^[a-z0-9_]+$", name):
+            raise ValueError("Plugin name must consist of lowercase letters, numbers, and underscores only")
+
+        self.name: str = name
+        self.components: dict[Component, Callable | Type] = {}
+
+        if name in TENSORFIELDS:
+            raise ValueError(f"Plugin '{name}' already registered")
+
+        TENSORFIELDS[name] = self
+
+    def register(self, obj: Type | Callable) -> Type | Callable:
+        if not hasattr(obj, "__name__"):
+            raise NameError(f"Object {obj} does not have a name")
+
+        name: str = str(obj.__name__)
+
+        if name in self.components:
+            raise ValueError(f"Component '{name}' already registered in plugin '{self.name}'")
+
+        match name:
+            case Component.Request:
+                if not isinstance(obj, type):
+                    raise TypeError("Request must be a class type")
+
+                if not issubclass(obj, Node):
+                    raise TypeError("Request must be a subclass of Node")
+
+                # for attr in Leaf.__annotations__.keys():
+                #     if not hasattr(obj, attr):
+                #         raise AttributeError(f"Request class must have a '{attr}' attribute")
+
+                # if getattr(obj, "type") != self.name:
+                #     raise ValueError(
+                #         f"Request class 'type' attribute must be '{self.name}', got '{getattr(obj, 'type')}'"
+                #     )
+
+            case Component.TensorField:
+                if not isinstance(obj, type):
+                    raise TypeError("TensorField must be a class type")
+
+                if not issubclass(obj, TensorFieldBase):
+                    raise TypeError("TensorField must be a subclass of TensorFieldBase")
+
+            case Component.Embedder:
+                if not isinstance(obj, type):
+                    raise TypeError("Embedder must be a class type")
+
+                if not issubclass(obj, EmbedderBase):
+                    raise TypeError("Embedder must be a subclass of EmbedderBase")
+
+                # confirm the init method is expecting structure and address
+                init_params = list(obj.__init__.__annotations__.keys())
+                if "structure" not in init_params or "address" not in init_params:
+                    raise TypeError("Embedder __init__ method must accept 'structure' and 'address' parameters")
+
+            case Component.Decoder:
+                if not isinstance(obj, type):
+                    raise TypeError("Decoder must be a class type")
+
+                if not issubclass(obj, DecoderBase):
+                    raise TypeError("Decoder must be a subclass of DecoderBase")
+
+                init_params = list(obj.__init__.__annotations__.keys())
+                if "structure" not in init_params or "address" not in init_params:
+                    raise TypeError("Decoder __init__ method must accept 'structure' and 'address' parameters")
+
+            case Component.loss:
+                if not callable(obj):
+                    raise TypeError("Loss must be a callable function")
+
+                expected_params: list[str] = ["module", "prediction", "batch", "strata"]
+                func_params: list[str] = list(obj.__annotations__.keys())
+
+                if not set(expected_params).issubset(set(func_params)):
+                    raise TypeError(
+                        f"Loss function must accept the following parameters: {expected_params}, got {func_params}"
+                    )
+
+            case Component.write:
+                if not callable(obj):
+                    raise TypeError("Write must be a callable function")
+
+                # check the signature of the function
+                expected_params: list[str] = ["module", "prediction"]
+                func_params: list[str] = list(obj.__annotations__.keys())
+
+                if func_params != expected_params:
+                    raise TypeError(
+                        f"Write function must accept the following parameters: {expected_params}, got {func_params}"
+                    )
+
+        self.components[name] = obj
+
+        return obj
+
+    @functools.cache
+    def __getattr__(self, key: Component) -> Callable | Type:
+        if key not in Component:
+            raise ValueError(f"Component '{key}' is not a valid Component enum value")
+
+        if key in self.components:
+            return self.components[key]
+
+        raise AttributeError(f"Plugin '{self.name}' has no component '{key}'")

json2vec/tensorfields/extensions/category.py

@@ -0,0 +1,484 @@
+from __future__ import annotations
+
+from functools import partial
+from multiprocessing import Manager
+from multiprocessing.managers import ListProxy, SyncManager
+from multiprocessing.synchronize import Lock
+from typing import TYPE_CHECKING, Annotated, Literal
+
+import numpy as np
+import pydantic
+import torch
+from beartype import beartype
+from ordered_set import OrderedSet
+from tensordict import TensorDict, tensorclass
+
+from json2vec.architecture.counter import Counter
+from json2vec.data.processing import apply, pad
+from json2vec.structs.enums import Metric, Strata, TensorKey, Tokens
+from json2vec.structs.packages import Parcel, Prediction
+from json2vec.structs.tree import Address
+from json2vec.tensorfields.base import (
+    DecoderBase,
+    EmbedderBase,
+    Plugin,
+    RequestBase,
+    TensorFieldBase,
+)
+
+if TYPE_CHECKING:
+    from json2vec.architecture.root import JSON2Vec
+    from json2vec.structs.experiment import Session, Structure
+
+category: Plugin = Plugin(name="category")
+
+class Vocabulary:
+    
+    def __init__(self, master: ListProxy, lock: Lock):
+
+        self.master: ListProxy[str] = master
+        self.lock: Lock = lock
+        self.vocab: OrderedSet[str] = OrderedSet(list(master))
+
+    def __call__(self, word: str, update: bool) -> int|None:
+
+        if word is None:
+            return None
+
+        if word in self.vocab:
+            return self.vocab.index(word)
+
+        if not update:
+            return None
+
+        # OK, it is not known locally... We will lock the global state and update the local vocab
+        with self.lock:
+            self.vocab: OrderedSet[str] = OrderedSet(list(self.master))
+
+            if word not in self.vocab:
+                self.vocab.add(word)
+                self.master.append(word)
+
+        return self.vocab.index(word)
+
+    def __len__(self) -> int:
+
+        return len(self.vocab)
+
+class OnlineVocabularyModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.manager: SyncManager = Manager()
+        self.master: ListProxy[str] = self.manager.list()
+        self.lock: Lock = self.manager.Lock()
+        self._snapshot_cache: list[str] | None = None
+        self._snapshot_size: int = -1
+
+    def _save_to_state_dict(self, state_dict, prefix, keep_vars):
+        super()._save_to_state_dict(state_dict, prefix, keep_vars)
+
+        state_dict[prefix + "vocabulary"] = list(self.master)
+
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata,
+        strict, missing_keys, unexpected_keys, error_msgs
+    ):
+
+        vocab: list[str] = state_dict.pop(prefix + "vocabulary")
+        self.master: ListProxy[str] = self.manager.list(vocab)
+        self._snapshot_cache = None
+        self._snapshot_size = -1
+
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata,
+            strict, missing_keys, unexpected_keys, error_msgs
+        )
+    
+    @property
+    def state(self) -> Vocabulary:
+
+        return Vocabulary(master=self.master, lock=self.lock)
+
+    def snapshot(self) -> list[str]:
+        size = len(self.master)
+        if self._snapshot_cache is None or self._snapshot_size != size:
+            self._snapshot_cache = list(self.master)
+            self._snapshot_size = size
+
+        return self._snapshot_cache
+
+
+@category.register
+class Request(RequestBase):
+    type: Literal["category"]
+    max_vocab_size: Annotated[int, pydantic.Field(gt=0, default=10_000)]
+    n_bands: Annotated[int, pydantic.Field(gt=0, default=8)]
+    topk: Annotated[list[int], pydantic.Field(default_factory=list)]
+
+    @pydantic.model_validator(mode="after")
+    def check_topk(self):
+
+        # enforce uniqueness
+        self.topk = sorted(set(self.topk))
+
+        for topk in self.topk:
+            if not isinstance(topk, int):
+                raise ValueError("topk values must be integers")
+
+            if topk <= 0:
+                raise ValueError("topk values must be positive")
+
+            if topk == 1:
+                raise ValueError("topk values must not be 1")
+
+            if topk >= self.max_vocab_size:
+                raise ValueError("topk values must be less than max_vocab_size")
+
+        return self
+
+    
+
+@category.register
+@tensorclass
+class TensorField(TensorFieldBase):
+    state: torch.Tensor
+    content: torch.Tensor
+    trainable: torch.Tensor
+    targets: TensorDict[TensorKey, torch.Tensor]
+
+    @classmethod
+    def new(
+        cls,
+        values: list,
+        address: Address,
+        session: Session,
+        strata: Strata,
+        state: Vocabulary,
+    ) -> TensorFieldBase:
+
+        context_shape: tuple[int, ...] = session.structure.shapes[address]
+
+        tokens = apply(values, partial(state, update=(strata == Strata.train)))
+
+        if len(state) > (max_vocab_size := session.structure.requests[address].max_vocab_size):
+            print(f"Vocab in address {address} exceeds max vocab size of {max_vocab_size}")
+
+        data, states = pad(
+            nested=tokens,
+            shape=(len(values), *context_shape),
+            dtype=np.int64,
+            pad_value=0,
+        )
+
+        state_tensor = torch.tensor(states, dtype=torch.int64)
+        content = torch.tensor(data=data, dtype=torch.int64)
+
+        return cls(
+            state=state_tensor,
+            content=content,
+            trainable=torch.zeros_like(input=state_tensor, dtype=torch.bool),
+            targets=TensorDict({}),
+            batch_size=len(values),
+        )
+
+    def mask(self, p_mask: float):
+        mask_token = torch.full_like(input=self.state, fill_value=Tokens.masked.value)
+        is_masked = torch.rand_like(input=self.state, dtype=torch.float).lt(other=p_mask)
+
+        if TensorKey.state not in self.targets.keys():
+            self.targets[TensorKey.state] = self.state.clone()
+
+        if TensorKey.content not in self.targets.keys():
+            self.targets[TensorKey.content] = self.content.clone()
+
+        self.state = self.state.masked_scatter(is_masked, mask_token)
+        self.content = self.content.masked_scatter(is_masked, torch.zeros_like(input=self.content))
+
+        self.trainable |= is_masked
+
+    def prune(self, p_prune: float = 1.0):
+        prune_tokens = torch.full_like(input=self.state, fill_value=Tokens.pruned.value)
+
+        is_pruned = (
+            torch.rand(self.state.size(0), *([1] * (len(self.state.shape) - 1)), device=self.state.device)
+            .lt(p_prune)
+            .expand_as(self.state)
+        )
+
+        if TensorKey.state not in self.targets.keys():
+            self.targets[TensorKey.state] = self.state.clone()
+
+        if TensorKey.content not in self.targets.keys():
+            self.targets[TensorKey.content] = self.content.clone()
+
+        self.state = self.state.masked_scatter(is_pruned, prune_tokens)
+        self.content = self.content.masked_scatter(is_pruned, torch.zeros_like(input=self.content))
+
+        self.trainable |= is_pruned
+
+    @classmethod
+    def empty(
+        cls,
+        batch_size: int,
+        address: Address,
+        structure: Structure,
+    ):
+        shape: tuple[int, ...] = (batch_size, *structure.shapes[address])
+
+        state = torch.full(shape, Tokens.pruned)
+        content = torch.zeros(shape, dtype=torch.int64)
+
+        return cls(
+            state=state,
+            content=content,
+            trainable=torch.zeros_like(input=state, dtype=torch.bool),
+            targets=TensorDict({}),
+            batch_size=batch_size,
+        )
+
+
+@category.register
+class Embedder(EmbedderBase):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__(structure=structure, address=address)
+
+        request: Request = structure.requests[address]
+        self.origin: Address = address
+        self.destination: Address = request.parent.address
+        self.max_vocab_size: int = request.max_vocab_size
+
+        self.vocab: OnlineVocabularyModel = OnlineVocabularyModel()
+
+        self.embeddings = torch.nn.ModuleDict(
+            {
+                TensorKey.state.name: torch.nn.Embedding(
+                    num_embeddings=len(Tokens),
+                    embedding_dim=structure.d_model,
+                ),
+                TensorKey.content.name: torch.nn.Embedding(
+                    num_embeddings=request.max_vocab_size,
+                    embedding_dim=structure.d_model,
+                ),
+            }
+        )
+
+    @beartype
+    def forward(self, inputs: TensorFieldBase) -> Parcel:
+        N: int
+        dims: tuple[int, ...]
+
+        N, *dims = inputs.state.shape
+        state = inputs.state.reshape(-1)
+        content = inputs.content.reshape(-1)
+        valued = state.eq(Tokens.valued.value)
+
+        if valued.any() and (content.masked_select(valued) >= self.max_vocab_size).any().item():
+            raise ValueError(f"Token in address {self.origin} exceeds max vocab size of {self.max_vocab_size}")
+
+        safe_content = content.masked_fill(~valued, 0)
+
+        embeddings: torch.Tensor = (
+            self.embeddings[TensorKey.state.name](state) +
+            self.embeddings[TensorKey.content.name](safe_content) * valued.unsqueeze(-1)
+        ).reshape(N, *dims, -1)
+
+
+        return Parcel(
+            payload=embeddings,
+            origin=self.origin,
+            destination=self.destination,
+            batch_size=N,
+        )
+
+    @property
+    def state(self) -> Vocabulary:
+        return self.vocab.state
+
+
+
+@category.register
+class Decoder(DecoderBase):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__(structure=structure, address=address)
+
+        request: RequestBase = structure.requests[address]
+
+        self.linears = torch.nn.ModuleDict(
+            {
+                TensorKey.state.name: torch.nn.Linear(
+                    in_features=structure.d_model,
+                    out_features=len(Tokens),
+                ),
+                TensorKey.content.name: torch.nn.Linear(
+                    in_features=structure.d_model,
+                    out_features=request.max_vocab_size,
+                ),
+            }
+        )
+
+        self.counters = torch.nn.ModuleDict(
+            {
+                TensorKey.state.name: Counter(address=address, size=len(Tokens)),
+                TensorKey.content.name: Counter(address=address, size=request.max_vocab_size),
+            }
+        )
+
+    @beartype
+    def decode(self, pooled: torch.Tensor) -> TensorDict[TensorKey, torch.Tensor]:
+        return TensorDict(
+            source={
+                TensorKey.state: self.linears[TensorKey.state.name](pooled),
+                TensorKey.content: self.linears[TensorKey.content.name](pooled),
+            }
+        )
+
+
+@category.register
+def loss(
+    module: JSON2Vec,
+    prediction: Prediction,
+    batch: TensorFieldBase,
+    strata: Strata,
+) -> torch.Tensor:
+    decoder: Decoder = module.nodes[prediction.address].decoder
+    N: int = batch.targets[TensorKey.state].numel()
+    trainable = batch.trainable.reshape(N)
+
+    state_inputs = prediction.payload[TensorKey.state].reshape(N, -1)
+    state_targets = batch.targets[TensorKey.state].reshape(N)
+    decoder.counters[TensorKey.state.name](batch.targets[TensorKey.state])
+
+    loss: torch.Tensor = module.track(
+        (prediction.address, strata, Metric.loss, TensorKey.state),
+        value=(
+            torch.nn.functional.cross_entropy(
+                input=state_inputs,
+                target=state_targets,
+                weight=decoder.counters[TensorKey.state.name].weight,
+                reduction="none",
+            )
+            .masked_select(trainable)
+            .mean()
+        )
+    )
+
+    module.track(
+        (prediction.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 loss
+
+    content_inputs = prediction.payload[TensorKey.content].reshape(N, -1)
+    content_targets = batch.targets[TensorKey.content].reshape(N)
+    content_counter_values = content_targets.masked_select(state_targets.eq(Tokens.valued.value))
+    if content_counter_values.numel() > 0:
+        decoder.counters[TensorKey.content.name](content_counter_values)
+
+    loss += module.track(
+        (prediction.address, strata, Metric.loss, TensorKey.content),
+        value=(
+            torch.nn.functional.cross_entropy(
+                input=content_inputs,
+                target=content_targets,
+                weight=decoder.counters[TensorKey.content.name].weight,
+                reduction="none",
+            )
+            .masked_select(valued)
+            .mean()
+        )
+    )
+
+    for topk in module.session.structure.requests[prediction.address].topk:
+        module.track(
+            (prediction.address, strata, Metric.accuracy, f"top{topk}"),
+            value=(
+                content_inputs
+                .topk(k=topk, dim=1)
+                .indices.eq(content_targets.unsqueeze(1))
+                .any(dim=1)
+                .masked_select(valued).float().mean()
+            )
+        )
+
+    module.track(
+        (prediction.address, strata, Metric.accuracy, TensorKey.content),
+        value=content_inputs.argmax(dim=1).eq(content_targets).masked_select(valued).float().mean(),
+    )
+
+    return loss
+
+
+@category.register
+def write(module: JSON2Vec, prediction: Prediction):
+
+    node = module.nodes[prediction.address]
+    state_logits: torch.Tensor = prediction.payload[TensorKey.state]
+    content_logits: torch.Tensor = prediction.payload[TensorKey.content]
+
+    tokens = 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())
+    }
+
+    vocab = np.array(node.embedder.vocab.snapshot(), dtype=object)
+    content_shape = tuple(state_distribution.shape[:-1])
+    content_labels = np.full(content_shape, None, dtype=object)
+    content_probabilities = np.zeros(content_shape, dtype=np.float32)
+
+    requested_ks: list[int] = module.session.structure.requests[prediction.address].topk
+    max_requested_k: int = max(requested_ks, default=0)
+
+    def _pack_candidates(labels: np.ndarray, probabilities: np.ndarray) -> list[dict[str, float]] | list:
+        if labels.ndim == 1:
+            return [
+                {"label": str(label), "probability": float(probability)}
+                for label, probability in zip(labels.tolist(), probabilities.tolist())
+            ]
+
+        return [_pack_candidates(labels[index], probabilities[index]) for index in range(labels.shape[0])]
+
+    def _empty_candidates(shape: tuple[int, ...]) -> list | None:
+        if len(shape) == 0:
+            return []
+
+        return [_empty_candidates(shape[1:]) for _ in range(shape[0])]
+
+    topk_payload: list | None = _empty_candidates(content_shape)
+
+    if len(vocab) > 0:
+        narrow: torch.Tensor = content_logits.narrow(dim=-1, start=0, length=len(vocab))
+        log_norm = narrow.logsumexp(dim=-1, keepdim=True)
+        max_logits, max_indices = narrow.max(dim=-1)
+        content_probabilities = (max_logits - log_norm.squeeze(-1)).exp().detach().float().cpu().numpy()
+
+        max_indices_np: np.ndarray = max_indices.detach().cpu().numpy().astype(np.int32)
+        content_labels = vocab[max_indices_np]
+
+        if max_requested_k > 0:
+            topk: int = min(max_requested_k, narrow.shape[-1])
+            topk_logits, topk_indices = narrow.topk(k=topk, dim=-1)
+            topk_probabilities = (topk_logits - log_norm).exp()
+
+            topk_indices_np: np.ndarray = topk_indices.detach().cpu().numpy().astype(np.int32)
+            topk_labels_np: np.ndarray = vocab[topk_indices_np]
+            topk_probabilities_np: np.ndarray = topk_probabilities.detach().float().cpu().numpy()
+            topk_payload = _pack_candidates(
+                labels=topk_labels_np,
+                probabilities=topk_probabilities_np,
+            )
+
+    return {
+        TensorKey.state.name: state_payload,
+        TensorKey.content.name: {
+            TensorKey.value.name: content_labels,
+            TensorKey.probability.name: content_probabilities,
+            TensorKey.topk.name: topk_payload,
+        },
+    }