json2vec 0.1.0__py3-none-any.whl

json2vec/__main__.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import argparse
+
+from loguru import logger
+
+from json2vec.entrypoints import execute
+from json2vec.structs.experiment import Experiment
+
+
+def train() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment", type=str, default=None)
+    parser.add_argument("--name", type=str, default=None)
+    parser.add_argument("--notes", type=str, default=None)
+    parser.add_argument("--experiments", type=str, default="experiments")
+    args = parser.parse_args()
+
+    experiment: Experiment = Experiment.from_config(
+        args.experiments,
+        experiment=args.experiment,
+        name=args.name,
+        notes=args.notes,
+    )
+    outputs = execute(experiment=experiment)
+
+    for session_name, output in outputs.items():
+        logger.info(f"session={session_name} output={output}")
+
+
+if __name__ == "__main__":
+    train()

json2vec/architecture/attention.py

@@ -0,0 +1,64 @@
+import math
+
+import torch
+
+from json2vec.architecture.rotary import RotaryEmbedding
+
+
+class RotaryMultiheadAttention(torch.nn.Module):
+    def __init__(self, d_model: int, nhead: int, dropout: float):
+        super().__init__()
+
+        if d_model % nhead != 0:
+            raise ValueError("d_model must be divisible by nhead")
+
+        self.d_model = d_model
+        self.nhead = nhead
+        self.head_dim = d_model // nhead
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+
+        self.q_proj = torch.nn.Linear(d_model, d_model)
+        self.k_proj = torch.nn.Linear(d_model, d_model)
+        self.v_proj = torch.nn.Linear(d_model, d_model)
+        self.out_proj = torch.nn.Linear(d_model, d_model)
+
+        self.rotary = RotaryEmbedding(d_model=self.head_dim)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+    def splitheads(self, inputs: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = inputs.shape
+        return inputs.reshape(batch, seq_len, self.nhead, self.head_dim).transpose(1, 2)
+
+    def rotate(self, inputs: torch.Tensor) -> torch.Tensor:
+        batch, nhead, seq_len, head_dim = inputs.shape
+        rotated = self.rotary(inputs.reshape(batch * nhead, seq_len, head_dim))
+        return rotated.reshape(batch, nhead, seq_len, head_dim)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_padding_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        q = self.rotate(self.splitheads(self.q_proj(query)))
+        k = self.rotate(self.splitheads(self.k_proj(key)))
+        v = self.splitheads(self.v_proj(value))
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        if key_padding_mask is not None:
+            mask = key_padding_mask
+            all_masked = mask.all(dim=1)
+            if all_masked.any():
+                mask = mask.clone()
+                mask[all_masked, 0] = False
+            scores = scores.masked_fill(mask[:, None, None, :], torch.finfo(scores.dtype).min)
+
+        probs = torch.softmax(scores, dim=-1)
+        probs = self.dropout(probs)
+
+        context = torch.matmul(probs, v)
+        context = context.transpose(1, 2).reshape(query.shape[0], query.shape[1], self.d_model)
+
+        return self.out_proj(context)

json2vec/architecture/counter.py

@@ -0,0 +1,37 @@
+import torch
+
+from json2vec.structs.tree import Address
+
+
+class Counter(torch.nn.Module):
+    def __init__(self, address: Address, size: int):
+        super().__init__()
+
+        self.size: int = size
+
+        # init with ones to avoid division by zero
+        # it doesn't matter much since we will normalize over time
+        self.register_buffer("counts", torch.ones(size, dtype=torch.int64))
+        self.is_full: bool = False
+
+    @torch.no_grad()
+    def forward(self, values: torch.Tensor):
+        if self.training and not self.is_full:
+            next_count_max = int(self.counts.max().item()) + int(values.numel())
+            could_overflow = next_count_max > torch.iinfo(self.counts.dtype).max
+
+            if could_overflow:
+                # if we are approaching the max value, we stop counting and assume the counts are full
+                self.is_full = True
+                return values
+
+            self.counts += torch.bincount(values.view(-1), minlength=self.counts.shape[0]).to(self.counts.dtype)
+
+        return values
+
+    @property
+    @torch.no_grad()
+    def weight(self) -> torch.Tensor:
+        counts = self.counts.to(dtype=torch.float32)
+        weights = counts.rsqrt()
+        return weights * (counts.sum() / (weights * counts).sum())

json2vec/architecture/encoder.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from json2vec.architecture.attention import RotaryMultiheadAttention
+from json2vec.architecture.pool import LearnedQueryCrossAttention
+from json2vec.structs.packages import Parcel
+from json2vec.structs.tree import Address
+
+if TYPE_CHECKING:
+    from json2vec.structs.structure import Structure
+
+
+class RotaryTransformerEncoderLayer(torch.nn.Module):
+    def __init__(self, d_model: int, nhead: int, dropout: float, ffn_multiplier: int = 4):
+        super().__init__()
+
+        self.attention_norm = torch.nn.LayerNorm(normalized_shape=d_model)
+        self.ffn_norm = torch.nn.LayerNorm(normalized_shape=d_model)
+
+        self.attention = RotaryMultiheadAttention(d_model=d_model, nhead=nhead, dropout=dropout)
+
+        hidden = d_model * ffn_multiplier
+        self.ffn = torch.nn.Sequential(
+            torch.nn.Linear(in_features=d_model, out_features=hidden),
+            torch.nn.GELU(),
+            torch.nn.Dropout(p=dropout),
+            torch.nn.Linear(in_features=hidden, out_features=d_model),
+            torch.nn.Dropout(p=dropout),
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        normed = self.attention_norm(inputs)
+        inputs = inputs + self.attention(normed, normed, normed)
+        return inputs + self.ffn(self.ffn_norm(inputs))
+
+
+class ContextEncoder(torch.nn.Module):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__()
+
+        context = structure.contexts[address]
+
+        self.origin: Address = address
+        self.destination: Address = context.parent.address
+
+        layers: list[RotaryTransformerEncoderLayer] = []
+        for _ in range(context.n_layers):
+            layers.append(
+                RotaryTransformerEncoderLayer(
+                    d_model=structure.d_model,
+                    nhead=context.n_heads,
+                    dropout=structure.dropout,
+                )
+            )
+
+        self.encoder = torch.nn.ModuleList(layers)
+
+        self.pool = LearnedQueryCrossAttention(
+            n_context=context.n_outputs,
+            d_model=structure.d_model,
+            nhead=context.n_heads,
+            dropout=structure.dropout,
+            n_linear=context.n_linear,
+        )
+
+    def forward(self, parcels: list[Parcel]) -> Parcel:
+        payloads: list[torch.Tensor] = []
+        for parcel in parcels:
+            payloads.append(parcel.payload)
+
+        concatenated: torch.Tensor = torch.cat(payloads, dim=-2)
+        N, *dims, L, C = concatenated.shape
+        encoded: torch.Tensor = concatenated.reshape(-1, L, C)
+
+        for layer in self.encoder:
+            encoded = layer(encoded)
+
+        pooled: torch.Tensor = self.pool(encoded).reshape(N, *dims[:-1], -1, C)
+
+        return Parcel(
+            payload=pooled,
+            origin=self.origin,
+            destination=self.destination,
+            batch_size=N,
+        )

json2vec/architecture/node.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from json2vec.architecture.encoder import ContextEncoder
+from json2vec.structs.tree import Address, Node
+from json2vec.tensorfields.base import (
+    TENSORFIELDS,
+    DecoderBase,
+    EmbedderBase,
+    Plugin,
+)
+
+if TYPE_CHECKING:
+    from json2vec.structs.config import Structure
+
+
+class NodeModule(torch.nn.Module):
+    def __init__(self, structure: Structure, address: Address):
+        super().__init__()
+
+        if address in structure.requests:
+            request: Node = structure.requests[address]
+            plugin: Plugin = TENSORFIELDS[request.type]
+            self.embedder: EmbedderBase = plugin.Embedder(structure=structure, address=address)
+            self.decoder: DecoderBase = plugin.Decoder(structure=structure, address=address)
+
+        elif address in structure.contexts:
+            self.encoder: ContextEncoder = ContextEncoder(structure=structure, address=address)
+
+        else:
+            raise ValueError("how did we get here?")

json2vec/architecture/pool.py

@@ -0,0 +1,61 @@
+import torch
+
+from json2vec.architecture.attention import RotaryMultiheadAttention
+
+
+class CrossAttentionBlock(torch.nn.Module):
+    def __init__(self, d_model: int, nhead: int, dropout: float, ffn_multiplier: int):
+        super().__init__()
+
+        self.attention_norm = torch.nn.LayerNorm(normalized_shape=d_model)
+        self.ffn_norm = torch.nn.LayerNorm(normalized_shape=d_model)
+        self.attention = RotaryMultiheadAttention(d_model=d_model, nhead=nhead, dropout=dropout)
+
+        hidden = d_model * ffn_multiplier
+        self.ffn = torch.nn.Sequential(
+            torch.nn.Linear(in_features=d_model, out_features=hidden),
+            torch.nn.GELU(),
+            torch.nn.Dropout(p=dropout),
+            torch.nn.Linear(in_features=hidden, out_features=d_model),
+            torch.nn.Dropout(p=dropout),
+        )
+
+    def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
+        attended = self.attention(self.attention_norm(queries), memory, memory)
+        queries = queries + attended
+        return queries + self.ffn(self.ffn_norm(queries))
+
+
+class LearnedQueryCrossAttention(torch.nn.Module):
+    def __init__(
+        self,
+        n_context: int,
+        d_model: int,
+        nhead: int,
+        dropout: float,
+        n_linear: int = 1,
+        ffn_multiplier: int = 4,
+    ):
+        super().__init__()
+
+        self.queries = torch.nn.Parameter(torch.normal(mean=0.0, std=1e-2, size=(n_context, d_model)))
+        self.blocks = torch.nn.ModuleList()
+        for _ in range(n_linear):
+            self.blocks.append(
+                CrossAttentionBlock(
+                    d_model=d_model,
+                    nhead=nhead,
+                    dropout=dropout,
+                    ffn_multiplier=ffn_multiplier,
+                )
+            )
+        self.norm = torch.nn.LayerNorm(normalized_shape=d_model)
+
+    def forward(self, memory: torch.Tensor) -> torch.Tensor:
+        N, _, _ = memory.shape
+        queries = self.queries.unsqueeze(0).expand(N, -1, -1)
+
+        for block in self.blocks:
+            queries = block(queries=queries, memory=memory)
+
+        return self.norm(queries)

json2vec/architecture/root.py

@@ -0,0 +1,338 @@
+import math
+import traceback
+from collections import defaultdict
+from functools import cache, partialmethod, wraps
+from typing import Any, NotRequired, Self, TypedDict
+
+import lightning.pytorch as lit
+import torch
+from beartype import beartype
+from loguru import logger
+from tensordict import TensorDict
+
+from json2vec.architecture.encoder import ContextEncoder
+from json2vec.architecture.node import NodeModule
+from json2vec.data.datasets import dataloader, mock
+from json2vec.structs.enums import Metric, Strata, TensorKey
+from json2vec.structs.experiment import Session
+from json2vec.structs.packages import Embedding, Parcel, Prediction
+from json2vec.structs.tree import Address
+from json2vec.tensorfields.base import (
+    TENSORFIELDS,
+    DecoderBase,
+    EmbedderBase,
+    Plugin,
+    RequestBase,
+    TensorFieldBase,
+)
+
+
+class Output(TypedDict):
+    loss: NotRequired[torch.Tensor]
+    predictions: list[Prediction]
+
+
+@beartype
+def step(
+    module: "JSON2Vec",
+    batch: TensorDict[Address, TensorFieldBase],
+    batch_idx: int,
+    strata: Strata,
+) -> Output:
+    predictions: list[Prediction] = module.forward(batch)
+
+    if strata == Strata.predict:
+        return Output(predictions=predictions)
+
+    losses: list[torch.Tensor] = []
+
+    for prediction in predictions:
+
+        if isinstance(prediction, Embedding):
+            continue
+
+        address: Address = prediction.address
+        request: RequestBase = module.session.structure.requests[address]
+        extension: Plugin = TENSORFIELDS[request.type]
+
+        loss: torch.Tensor = extension.loss(module=module, prediction=prediction, batch=batch[address], strata=strata)
+        losses.append(loss * torch.tensor(request.weight))
+
+    if len(losses) == 0:
+        # under idealistic circumstances this would never happen.
+        # but with small mask rates, batch sizes, and flat input data it is possible
+        logger.warning("no trainable fields in batch, returning zero loss")
+        loss: torch.Tensor = torch.tensor(0.0, device=batch.device, requires_grad=True)
+        return Output(loss=loss, predictions=[])
+
+    loss: torch.Tensor = module.track((Metric.loss, strata), value=torch.stack(losses).sum())
+
+    return Output(loss=loss, predictions=predictions)
+
+
+def compile(fn):
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        model = fn(*args, **kwargs)
+        try:
+            import thunder
+
+            model = thunder.compile(model)
+            logger.info("successfully compiled module with thunder")
+        except Exception:
+            traceback.print_exc()
+            logger.info("[thunder] Returning uncompiled model instead.")
+        return model
+
+    return wrapper
+
+
+@cache
+def groupname(names: tuple[str, ...]) -> str:
+    assert len(names) > 1
+
+    group, *keys = tuple(map(lambda x: x.replace("/", ":").lower(), names))
+
+    key: str = ":".join(list(keys))
+
+    return f"{group}/{key}"
+
+
+class JSON2Vec(lit.LightningModule):
+    @beartype
+    def __init__(self, session: Session):
+
+        super().__init__()
+
+        self.session: Session = session
+
+        self.nodes: torch.nn.ModuleDict[str, NodeModule] = torch.nn.ModuleDict()
+
+        for address in self.session.structure.requests | self.session.structure.contexts:
+            self.nodes[address] = NodeModule(structure=self.session.structure, address=address)
+
+        self.example_input_array = mock(structure=session.structure)
+
+        logger.bind(
+            component="model",
+            session=self.session.name,
+            structure=self.session.structure.name,
+            requests=len(self.session.structure.requests),
+            contexts=len(self.session.structure.contexts),
+            outputs=len(self.session.output),
+        ).info("initialized JSON2Vec module")
+
+    def track(self, names: tuple[str, ...], /, value: torch.Tensor) -> torch.Tensor:
+        self.log(
+            name=groupname(names),
+            value=value,
+            on_step=False,
+            on_epoch=True,
+            sync_dist=True,
+            batch_size=self.session.structure.batch_size,
+        )
+
+        return value
+
+    @property
+    def state(self) -> dict[Address, Any]:
+        return {
+            address: node.embedder.state
+            for address, node in self.nodes.items()
+            if hasattr(node, "embedder") and hasattr(node.embedder, "state")
+        }
+
+    @beartype
+    def forward(self, inputs: TensorDict[Address, TensorFieldBase]) -> list[Prediction]:
+        processed: dict[Address, list[Parcel]] = defaultdict(list)
+        outgoing: dict[Address, Parcel] = {}
+        predictions: list[Prediction] = []
+
+        for address in self.session.structure.requests.keys():
+            if address in self.session.pruned:
+                continue
+
+            tensorfield: TensorFieldBase = inputs[address]
+            embedder: EmbedderBase = self.nodes[address].embedder
+            embedding: Parcel = embedder(tensorfield)
+            processed[embedding.destination].append(embedding)
+            outgoing[embedding.origin] = embedding
+
+            if address in self.session.output:
+                predictions.append(Embedding.from_parcel(embedding))
+
+        # DAG traversal from leaves to root
+        for depth in reversed(self.session.structure.depthwise):
+            # these are order-independent within the same depth
+            for address in depth:
+
+                if len(processed[address]) == 0:
+                    continue
+
+                encoder: ContextEncoder = self.nodes[address].encoder
+                encoding: Parcel = encoder(processed[address])
+                processed[encoding.destination].append(encoding)
+                outgoing[encoding.origin] = encoding
+
+                if address in self.session.output:
+                    predictions.append(Embedding.from_parcel(encoding))
+
+        for address in self.session.structure.requests.keys():
+
+            if (torch.any(inputs[address].trainable)) or (address in self.session.pruned):
+
+                heritage: list[Address] = self.session.structure.requests[address].heritage
+                parcels: list[Parcel] = [
+                    outgoing[address] for address in heritage
+                    if address not in self.session.pruned and address in outgoing.keys()
+                ]
+
+                decoder: DecoderBase = self.nodes[address].decoder
+                predictions.append(decoder(parcels))
+
+        return predictions
+
+    @beartype
+    def configure_optimizers(self):
+
+        if self.session.learning_rate is None:
+            raise ValueError("learning_rate must be defined for optimizer configuration")
+
+        class GroupedParameter(TypedDict):
+            params: list[torch.nn.Parameter]
+            weight_decay: float
+
+        params: dict[str, GroupedParameter] = dict(
+            with_decay = GroupedParameter(params=[], weight_decay=self.session.weight_decay),
+            no_decay = GroupedParameter(params=[], weight_decay=0.0),
+        )
+
+        for name, parameter in self.named_parameters():
+            if not parameter.requires_grad:
+                continue
+
+            if name.endswith(".bias") or parameter.ndim <= 1 or "norm" in name.lower():
+                params["no_decay"]["params"].append(parameter)
+            else:
+                params["with_decay"]["params"].append(parameter)
+
+
+        optimizer = torch.optim.AdamW(params=list(params.values()), lr=self.session.learning_rate, betas=(0.9, 0.95))
+        trainable_parameters: int = sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
+        logger.bind(
+            component="optimizer",
+            session=self.session.name,
+            learning_rate=self.session.learning_rate,
+            weight_decay=self.session.weight_decay,
+            trainable_parameters=trainable_parameters,
+            warmup_ratio=self.session.warmup_ratio,
+            min_lr_ratio=self.session.min_lr_ratio,
+        ).info("configured AdamW optimizer")
+
+        total = int(getattr(self.trainer, "estimated_stepping_batches", 0) or 0)
+
+        if total <= 0:
+            return optimizer
+
+        warmup = max(1, int(total * self.session.warmup_ratio))
+        min_lr_ratio = self.session.min_lr_ratio
+
+        def schedule(step: int) -> float:
+
+            if step < warmup:
+                return float(step + 1) / float(warmup)
+
+            ratio = float(step - warmup) / float(max(1, total - warmup))
+
+            progress = min(1.0, ratio)
+
+            cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+            return min_lr_ratio + (1.0 - min_lr_ratio) * cosine
+
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule)
+
+        return dict(
+            optimizer= optimizer,
+            lr_scheduler=dict(
+                scheduler = scheduler,
+                interval = "step",
+                frequency = 1,
+            ),
+        )
+
+    def on_save_checkpoint(self, checkpoint):
+        logger.bind(component="checkpoint", session=self.session.name).info("serializing session")
+        checkpoint["session"] = self.session.model_dump()
+
+    def on_load_checkpoint(self, checkpoint):
+        logger.bind(component="checkpoint").info("loading session from checkpoint payload")
+        if "session" in checkpoint and getattr(self, "session", None) is None:
+            self.session = Session.model_validate(checkpoint["session"])
+
+        if getattr(self, "session", None) is None:
+            raise ValueError("missing session in checkpoint and constructor")
+
+    @classmethod
+    def get_or_create(
+        cls,
+        session: Session|None = None,
+        checkpoint: str | None = None,
+    ) -> Self:
+
+        if checkpoint is None:
+            logger.bind(component="model_factory").info("creating new JSON2Vec model")
+            if session is None:
+                raise ValueError("session is required when checkpoint is not provided")
+
+            model: "JSON2Vec" = cls(session=session)
+
+            return model
+
+        else:
+            logger.bind(component="model_factory", checkpoint=checkpoint).info("loading JSON2Vec model from checkpoint")
+            state = torch.load(checkpoint, weights_only=False, map_location="cpu")
+
+            model: "JSON2Vec" = cls(
+                session=session or Session.model_validate(state["session"]),
+            )
+
+            model.load_state_dict(state_dict=state["state_dict"], strict=False)
+            logger.bind(component="model_factory", checkpoint=checkpoint).info("restored model state from checkpoint")
+
+            return model
+
+    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]] = {}
+
+        for prediction in predictions:
+
+            if isinstance(prediction, Embedding):
+
+                embeddings[prediction.address] = Embedding.write(prediction)
+
+                continue
+
+            request: RequestBase = self.session.structure.requests[prediction.address]
+
+            extension: Plugin = TENSORFIELDS[request.type]
+
+            scribed: dict[TensorKey, Any]|None = extension.write(module=self, prediction=prediction)
+
+            if scribed is not None:
+                supervised[prediction.address] = Prediction.serialize(scribed)
+
+
+
+        return supervised, embeddings
+
+    training_step = partialmethod(step, strata=Strata.train)
+    validation_step = partialmethod(step, strata=Strata.validate)
+    test_step = partialmethod(step, strata=Strata.test)
+    predict_step = partialmethod(step, strata=Strata.predict)
+
+    train_dataloader = partialmethod(dataloader, strata=Strata.train)
+    val_dataloader = partialmethod(dataloader, strata=Strata.validate)
+    test_dataloader = partialmethod(dataloader, strata=Strata.test)
+    predict_dataloader = partialmethod(dataloader, strata=Strata.predict)