VectorMesh

A PyTorch framework for embed once, reuse many times.

A large pretrained encoder — a BERT-family text model, a vision CNN/ViT, or a hand-written regex feature extractor — is run over a dataset exactly once, by whoever has the hardware and the judgement to pick it. The resulting vectors go to disk as a VectorCache: a versioned, documented artefact. Everything after that trains a small head on those frozen vectors, cheaply enough to run on a laptop CPU.

That split is the point. One embedding job, many models trained against its output — and the interesting design space (fusing representations, gating, mixtures of experts) lives entirely on the cheap side of it.

📚 Documentation

Full documentation is in docs/.

Installation

uv sync

Requires Python ≥ 3.12. See pyproject.toml for the full dependency list.

Quick start

from pathlib import Path

import torch
from torch.utils.data import DataLoader

from vectormesh import VectorCache
from vectormesh.components import FixedPadding, MaskedMeanAggregator, NeuralNet, Serial
from vectormesh.data import Collate, OneHot

cache = VectorCache.load(path=Path("artefacts/my_dataset_train"))
hidden_size = cache.metadata["legal_dutch"]["hidden_size"]

data = cache.dataset.map(OneHot(num_classes=32, label_col="labels", target_col="onehot"))
loader = DataLoader(
    data,
    batch_size=32,
    shuffle=True,
    collate_fn=Collate(
        embedding_col="legal_dutch",
        target_col="onehot",
        padder=FixedPadding(max_chunks=30),
    ),
)

pipeline = Serial([
    MaskedMeanAggregator(),                  # (batch, chunks, dim) -> (batch, dim)
    NeuralNet(hidden_size, out_size=32),     # (batch, dim)         -> (batch, 32)
])

Then hand pipeline and loader to mltrainer.Trainer — see Training for the full wiring.

Building a cache instead of loading one, extending a cache with extra feature columns, and the image path are all covered in The data layer.

What's in the box

Data — VectorCache, Vectorizer (chunked text), ImageVectorizer, RegexVectorizer, ChunkedRegexVectorizer, DatasetSchema, OneHot, Collate, CollateParallel, LabelEncoder.

Components — pipelines (Serial, Parallel), padding (FixedPadding, DynamicPadding), aggregation (Mean, MaskedMean, Attention, RNN), neural blocks (NeuralNet, Projection, Attention, TransformerBlock), connectors (Concatenate2D, Concatenate3D, Stack2D), gating (Skip, Gate, Highway, MoE), augmentation (GaussianNoise), metrics (Accuracy, F1Score, MAE, MASE).

Details and signatures: Components, API reference.

Runtime type checking

VectorMesh annotates every forward with jaxtyping shapes, checked at runtime by beartype:

@jaxtyped(typechecker=beartype)
def forward(self, tensors: Float[Tensor, "batch chunks dim"]) -> Float[Tensor, "batch dim"]:
    return tensors.mean(dim=1)

This catches the failure mode that matters most here: nn.Linear accepts any leading shape, so feeding it (batch, chunks, dim) when you meant (batch, dim) raises no error — it just trains a model that quietly answers the wrong question. A BeartypeCallHintParamViolation naming a 3D tensor where a 2D one was expected almost always means a missing aggregator.

How to read those errors: Tensor contracts.

Notebooks

Walkthrough and exercises: Teaching path.

Scripts

Batch counterparts to the notebooks, for real dataset sizes:

uv run python scripts/create_cache_aktes.py    # embeddings
uv run python scripts/add_chunked_regex.py     # + chunk-aligned regex column
uv run python scripts/train_moe_parallel.py    # train

Full list: Training §6.5.

Development

uv run pytest -m "not integration"   # unit tests, no model downloads
uv run pytest                        # everything, downloads HF models
uv run ruff check src tests
uv run ty check

Project structure

src/vectormesh/
├── types.py               # VectorMeshError, Cachable, BaseComponent, TensorInput
├── data/                  # vectorizers, cache, schema, collation
└── components/            # pipelines, padding, aggregation, neural,
                           # connectors, gating, augmentation, metrics
docs/                      # documentation (start at docs/README.md)
notebooks/                 # tutorials, in order
scripts/                   # batch pipelines
references/                # source papers
tests/