vectormesh 0.1.0__py3-none-any.whl

vectormesh/__init__.py

@@ -0,0 +1,28 @@
+import sys
+from importlib.metadata import version
+
+from loguru import logger
+
+from vectormesh.data import (
+    BaseVectorizer,
+    LabelEncoder,
+    RegexVectorizer,
+    VectorCache,
+    Vectorizer,
+    build,
+)
+
+__version__ = version("vectormesh")
+
+__all__ = [
+    "VectorCache",
+    "LabelEncoder",
+    "build",
+    "BaseVectorizer",
+    "Vectorizer",
+    "RegexVectorizer",
+]
+
+logger.remove()
+logger.add(sys.stderr, level="INFO")
+logger.add("logs/dataset.log", rotation="10 MB", level="DEBUG")

vectormesh/components/__init__.py

@@ -0,0 +1,33 @@
+"""VectorMesh components module."""
+
+from .aggregation import (
+    AttentionAggregator,
+    BaseAggregator,
+    MeanAggregator,
+    RNNAggregator,
+)
+from .connectors import Concatenate2D, Stack2D
+from .gating import Gate, Highway, MoE, Skip
+from .neural import Attention, NeuralNet, Projection
+from .padding import DynamicPadding, FixedPadding
+from .pipelines import Parallel, Serial
+
+__all__ = [
+    "AttentionAggregator",
+    "BaseAggregator",
+    "MeanAggregator",
+    "RNNAggregator",
+    "Concatenate2D",
+    "Stack2D",
+    "Gate",
+    "Highway",
+    "Skip",
+    "MoE",
+    "NeuralNet",
+    "Projection",
+    "Attention",
+    "DynamicPadding",
+    "FixedPadding",
+    "Parallel",
+    "Serial",
+]

vectormesh/components/aggregation.py

@@ -0,0 +1,80 @@
+from abc import abstractmethod
+
+import torch
+import torch.nn as nn
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+
+from vectormesh.types import BaseComponent
+
+
+class BaseAggregator(BaseComponent):
+    """Base class for aggregating 3D -> 2D tensors.
+    We use "forward" to be compatible with nn.Module
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    @abstractmethod
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, embeddings: Float[Tensor, "batch _ dim"]
+    ) -> Float[Tensor, "batch dim"]:
+        """Aggregate from (batch, chunks, dim) to (batch, dim)."""
+        ...
+
+
+class MeanAggregator(BaseAggregator):
+    """Aggregate by taking mean over chunks.
+    no learnable parameters.
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, embeddings: Float[Tensor, "batch _ dim"]
+    ) -> Float[Tensor, "batch dim"]:
+        """Mean over chunks dimension."""
+        return embeddings.mean(dim=1)
+
+
+class AttentionAggregator(BaseAggregator):
+    """Aggregate using learnable attention over chunks.
+    Because attention does not handle variable-length sequences,
+    we actually get (batch, chunks, dim) where chunks is fixed.
+    """
+
+    def __init__(self, hidden_size: int):
+        """initialize learnable parameters."""
+        super().__init__()
+        self.attention = nn.Linear(hidden_size, 1)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, embeddings: Float[Tensor, "batch _ dim"]
+    ) -> Float[Tensor, "batch dim"]:
+        # attention_weights: (batch, _, 1)
+        attention_weights = torch.softmax(self.attention(embeddings), dim=1)
+        return (embeddings * attention_weights).sum(dim=1)
+
+
+class RNNAggregator(BaseAggregator):
+    """Aggregate using RNN over chunks.
+    return final hidden state.
+    """
+
+    def __init__(self, hidden_size: int):
+        """initialize learnable parameters."""
+        super().__init__()
+        self.rnn = torch.nn.GRU(
+            input_size=hidden_size, hidden_size=hidden_size, batch_first=True
+        )
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, embeddings: Float[Tensor, "batch _ dim"]
+    ) -> Float[Tensor, "batch dim"]:
+        output, _ = self.rnn(embeddings)
+        return output[:, -1, :]
+        return output[:, -1, :]

vectormesh/components/connectors.py

@@ -0,0 +1,40 @@
+import torch
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+
+from vectormesh.types import BaseComponent
+
+
+class Concatenate2D(BaseComponent):
+    """Concatenate tuples from parallel branches at last dimension.
+
+    input: ((batch dim1), (batch dim2), ...)
+    output: (batch ndim)
+
+    where ndim = dim1 + dim2 + ...
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, tensors: tuple[Float[Tensor, "batch dim"], ...]
+    ) -> Float[Tensor, "batch ndim"]:
+        return torch.cat(tensors, dim=-1)
+
+
+class Stack2D(BaseComponent):
+    """Stack n tuples from parallel branches, default 1st dimension.
+
+    input : ((batch dim1), (batch dim1), ...)
+    output: (batch nstack dim1)
+
+    where nstack is the number of tensors in the tuple
+    """
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, tensors: tuple[Float[Tensor, "batch dim1"], ...]
+    ) -> Float[Tensor, "batch nstack dim1"]:
+        return torch.stack(tensors, dim=1)
+        return torch.stack(tensors, dim=1)
+        return torch.stack(tensors, dim=1)

vectormesh/components/gating.py

@@ -0,0 +1,132 @@
+"""Residual and gating components for skip connections and gated transformations."""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+
+from vectormesh.types import BaseComponent
+
+
+class Skip(BaseComponent):
+    """Residual skip connection: output = batchnorm(transform(x) + projection(x))
+    - transform is the pipeline we want to apply to the input
+    - in_size is the dimensionality of the input; we need this for the layernorm
+    - projection is an optional pipeline, eg a Linear(in_size, out_size) if the
+    transform changes the dimensionality.
+    """
+
+    transform: nn.Module
+    projection: Optional[nn.Module]
+    layernorm: nn.LayerNorm
+
+    def __init__(
+        self,
+        transform: nn.Module,
+        in_size: int,
+        projection: Optional[nn.Module] = None,
+    ):
+        super().__init__()
+        self.transform = transform
+        self.projection = projection
+        self.layernorm = nn.LayerNorm(in_size)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(self, x: Float[Tensor, "..."]) -> Float[Tensor, "..."]:
+        # pre-norm (instead of post-norm) improves stability
+        x = self.layernorm(x)
+        residual = self.projection(x) if self.projection else x
+        transformed = self.transform(x)
+        return transformed + residual
+
+
+class Gate(BaseComponent):
+    """Simple gating: output = sigmoid(W·x) * x"""
+
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.project = nn.Linear(hidden_size, hidden_size)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(self, x: Float[Tensor, "batch dim"]) -> Float[Tensor, "batch dim"]:
+        return F.sigmoid(self.project(x)) * x
+
+
+class Highway(BaseComponent):
+    """Highway network: G * T(x) + (1-G) * x"""
+
+    def __init__(self, transform: nn.Module, hidden_size: int):
+        super().__init__()
+        self.transform = transform
+        self.project = nn.Linear(hidden_size, hidden_size)
+        self.norm = nn.LayerNorm(hidden_size)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(self, x: Float[Tensor, "batch dim"]) -> Float[Tensor, "batch dim"]:
+        # pre-norm (instead of post-norm) improves stability
+        x = self.norm(x)
+        gate = F.sigmoid(self.project(x))
+        transformed = self.transform(x)
+        return gate * transformed + (1 - gate) * x
+
+
+class MoE(BaseComponent):
+    """
+    See https://arxiv.org/abs/1701.06538 for paper
+    """
+
+    def __init__(self, experts, hidden_size, out_size, top_k, noisy_gating=True):
+        super().__init__()
+        self.experts = nn.ModuleList(experts)
+        self.router = nn.Linear(hidden_size, len(experts))
+
+        self.w_noise = nn.Linear(hidden_size, len(experts))
+        self.noisy_gating = noisy_gating
+        self.top_k = top_k
+        self.num_experts = len(experts)
+        self.out_size = out_size
+
+    def forward(self, x):
+        clean_logits = self.router(x)
+
+        # self.training is automatically managed by .eval() and .train()
+        if self.noisy_gating and self.training:
+            raw_noise_stddev = self.w_noise(x)
+            noise_stddev = F.softplus(raw_noise_stddev) + 1e-2
+            noise = torch.randn_like(clean_logits) * noise_stddev
+            noisy_logits = clean_logits + noise
+        else:
+            noisy_logits = clean_logits
+
+        # We set non-top-k logits to -inf so Softmax drives them to absolute zero
+        top_logits, top_indices = noisy_logits.topk(self.top_k, dim=1)
+        full_logits = torch.full_like(noisy_logits, float("-inf"))
+        full_logits.scatter_(1, top_indices, top_logits)
+
+        router_probs = F.softmax(full_logits, dim=1)
+
+        final_output = torch.zeros(x.size(0), self.out_size, device=x.device)
+
+        for i in range(self.num_experts):
+            mask = (top_indices == i).any(dim=1)
+
+            if mask.any():
+                expert_input = x[mask]
+                expert_output = self.experts[i](expert_input)
+
+                expert_weights = router_probs[mask, i].unsqueeze(-1)
+
+                final_output[mask] += expert_output * expert_weights
+
+        # TODO: the paper implements importance loss
+        # to encourage balanced expert usage
+        #
+        # importance = router_probs.sum(dim=0)
+        # imp_loss = (importance.std() / (importance.mean() + 1e-10)).pow(2)
+        # return final_output, imp_loss
+
+        return final_output

vectormesh/components/metrics.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Iterator, Union
+
+import numpy as np
+import torch
+
+NumericType = Union[torch.Tensor, np.ndarray]
+
+
+class Metric(ABC):
+    """Base class for all metrics with proper device and type handling"""
+
+    @abstractmethod
+    def _compute(self, y: torch.Tensor, yhat: torch.Tensor) -> torch.Tensor:
+        """Actual metric computation - implemented by subclasses"""
+        pass
+
+    def __call__(self, y: NumericType, yhat: NumericType) -> float:
+        """Handle device/type conversion and return final metric"""
+        # Convert to tensors if needed
+        if isinstance(y, np.ndarray):
+            y = torch.from_numpy(y)
+        if isinstance(yhat, np.ndarray):
+            yhat = torch.from_numpy(yhat)
+
+        # Ensure we're working with the same device
+        device = y.device
+        y = y.to(device)
+        yhat = yhat.to(device)
+
+        # Compute metric
+        result = self._compute(y, yhat)
+
+        # Return as float for consistency
+        return float(result.cpu().detach())
+
+    @abstractmethod
+    def __repr__(self) -> str:
+        pass
+
+
+class MAE(Metric):
+    def _compute(self, y: torch.Tensor, yhat: torch.Tensor) -> torch.Tensor:
+        return torch.mean(torch.abs(y - yhat))
+
+    def __repr__(self) -> str:
+        return "MAE"
+
+
+class MASE(Metric):
+    def __init__(self, train: Iterator, horizon: int) -> None:
+        self.horizon = horizon
+        with torch.no_grad():
+            self.scale = self._calculate_scale(train)
+
+    def _calculate_scale(self, train: Iterator) -> torch.Tensor:
+        elist = []
+        streamer = train.stream()  # type: ignore
+        for _ in range(len(train)):  # type: ignore
+            x, y = next(iter(streamer))
+            yhat = self._naive_predict(x)
+            e = torch.mean(torch.abs(y - yhat))
+            elist.append(e)
+        return torch.mean(torch.stack(elist))
+
+    def _naive_predict(self, x: torch.Tensor) -> torch.Tensor:
+        return x[..., -self.horizon :, :].squeeze(-1)
+
+    def _compute(self, y: torch.Tensor, yhat: torch.Tensor) -> torch.Tensor:
+        mae = torch.mean(torch.abs(y - yhat))
+        return mae / self.scale
+
+    def __repr__(self) -> str:
+        return f"MASE(scale={self.scale:.3f})"
+
+
+class Accuracy(Metric):
+    def _compute(self, y: torch.Tensor, yhat: torch.Tensor) -> torch.Tensor:
+        predictions = yhat.argmax(dim=1)
+        return (predictions == y).float().mean()
+
+    def __repr__(self) -> str:
+        return "Accuracy"
+
+
+class F1Score(Metric):
+    def __init__(self, average="micro", threshold=0.5, epsilon=1e-7):
+        self.average = average
+        self.threshold = threshold
+        self.epsilon = epsilon
+
+    def _compute(self, y: torch.Tensor, yhat: torch.Tensor) -> torch.Tensor:
+        y_prob = torch.sigmoid(yhat)
+        y_pred = (y_prob > self.threshold).float()
+
+        tp = y_pred * y
+        fp = y_pred * (1 - y)
+        fn = (1 - y_pred) * y
+
+        if self.average == "micro":
+            tp_sum = tp.sum()
+            fp_sum = fp.sum()
+            fn_sum = fn.sum()
+            f1 = (2 * tp_sum) / (2 * tp_sum + fp_sum + fn_sum + self.epsilon)
+
+        elif self.average == "macro":
+            tp_sum = tp.sum(dim=0)
+            fp_sum = fp.sum(dim=0)
+            fn_sum = fn.sum(dim=0)
+            f1_per_class = (2 * tp_sum) / (2 * tp_sum + fp_sum + fn_sum + self.epsilon)
+            f1 = f1_per_class.mean()
+
+        else:
+            raise ValueError("Average must be 'micro' or 'macro'")
+
+        return f1
+
+    def __repr__(self) -> str:
+        return f"F1({self.average})"

vectormesh/components/neural.py

@@ -0,0 +1,62 @@
+import torch.nn as nn
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+
+from vectormesh.types import BaseComponent
+
+
+class NeuralNet(BaseComponent):
+    """Two-layer feedforward network with GELU activation."""
+
+    def __init__(self, hidden_size: int, out_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.out_size = out_size
+        self.fc1 = nn.Linear(hidden_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, out_size)
+        self.activation = nn.GELU()
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, x: Float[Tensor, "batch {self.hidden_size}"]
+    ) -> Float[Tensor, "batch {self.out_size}"]:
+        return self.fc2(self.activation(self.fc1(x)))
+
+
+class Projection(BaseComponent):
+    """Linear projection layer."""
+
+    def __init__(self, hidden_size: int, out_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.out_size = out_size
+        self.proj = nn.Linear(hidden_size, out_size)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, x: Float[Tensor, "batch {self.hidden_size}"]
+    ) -> Float[Tensor, "batch {self.out_size}"]:
+        return self.proj(x)
+
+
+class Attention(nn.Module):
+    """Multi-head self-attention using PyTorch's implementation."""
+
+    def __init__(self, hidden_size: int, num_heads: int = 8, dropout: float = 0.1):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.attn = nn.MultiheadAttention(
+            embed_dim=hidden_size,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+
+    @jaxtyped(typechecker=beartype)
+    def forward(
+        self, x: Float[Tensor, "batch seq {self.hidden_size}"]
+    ) -> Float[Tensor, "batch seq {self.hidden_size}"]:
+        # Self-attention: query, key, value all come from x
+        attn_output, _ = self.attn(x, x, x, need_weights=False)
+        return attn_output

vectormesh/components/padding.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+import torch.nn.functional as F
+from beartype import beartype
+from jaxtyping import Float, jaxtyped
+from torch import Tensor
+from torch.nn.utils.rnn import pad_sequence
+
+
+class DynamicPadding:
+    @jaxtyped(typechecker=beartype)
+    def __call__(self, embeddings: list[Any]) -> Float[Tensor, "batch _ dim"]:
+        """Pad sequences to the maximum length in the batch."""
+        return pad_sequence(
+            embeddings,
+            batch_first=True,
+        )
+
+
+class FixedPadding:
+    def __init__(self, max_chunks: int):
+        self.max_chunks = max_chunks
+
+    @jaxtyped(typechecker=beartype)
+    def __call__(
+        self, embeddings: list[Any]
+    ) -> Float[Tensor, "batch {self.max_chunks} dim"]:
+        padded = pad_sequence(embeddings, batch_first=True)
+
+        current = padded.shape[1]
+        if current < self.max_chunks:
+            return F.pad(padded, (0, 0, 0, self.max_chunks - current))
+        else:
+            return padded[:, : self.max_chunks, :]

vectormesh/components/pipelines.py

@@ -0,0 +1,41 @@
+import torch.nn as nn
+from beartype import beartype
+from jaxtyping import jaxtyped
+
+from vectormesh.types import TensorInput
+
+
+class Serial(nn.Module):
+    """Sequential composition - just runs components in order.
+    Tensor checking happens via jaxtyping decorators on each component.
+    """
+
+    components: nn.ModuleList
+
+    def __init__(self, components: list[nn.Module]):
+        super().__init__()
+        self.components = nn.ModuleList(components)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(self, tensors: TensorInput) -> TensorInput:
+        """Execute pipeline. Type checking via component decorators."""
+        result = tensors
+        for component in self.components:
+            result = component(result)
+        return result
+
+
+class Parallel(nn.Module):
+    """Parallel composition - runs branches independently and returns tuple.
+    All branches receive the same input.
+    """
+
+    branches: nn.ModuleList
+
+    def __init__(self, branches: list[nn.Module]):
+        super().__init__()
+        self.branches = nn.ModuleList(branches)
+
+    @jaxtyped(typechecker=beartype)
+    def forward(self, tensors: TensorInput) -> TensorInput:
+        return tuple(branch(t) for branch, t in zip(self.branches, tensors))

vectormesh/data/__init__.py

@@ -0,0 +1,16 @@
+"""VectorMesh data components."""
+
+from .cache import VectorCache
+from .dataset import Collate, LabelEncoder, OneHot, build
+from .vectorizers import BaseVectorizer, RegexVectorizer, Vectorizer
+
+__all__ = [
+    "VectorCache",
+    "LabelEncoder",
+    "OneHot",
+    "Collate",
+    "build",
+    "BaseVectorizer",
+    "Vectorizer",
+    "RegexVectorizer",
+]