JSON2Vec

json2vec is a schema-driven framework for predictive modeling over nested, structured records without flattening them into a fixed feature table first.

The schema becomes the encoder: leaf tensorfield plugins encode raw values, array nodes aggregate child embeddings with transformer layers, and datatype-specific decoders reconstruct masked, targeted, or supervised fields from the surrounding hierarchy.

This supports self-supervised pretraining, supervised targets, embeddings, and schema evolution in one model surface. Customer/account/transaction data, flight itineraries, order fulfillment events, clickstream sessions, and other nested records can use the same machinery while keeping proprietary data, schemas, and checkpoints private.

What Makes This Different

• Attributed embeddings. The model can emit embeddings from any configured field or array, not only from the root. That makes branch-level similarity and retrieval workflows possible without flattening the record.
• Extensible data types for predictive modeling. Masked values, targeted fields, and explicit supervised targets all flow through the same datatype-specific heads. A new tensorfield type brings its own embedding, decoding, loss, and writing logic, so the framework stays reusable as schemas grow.
• Schema evolution is a first-class workflow. Between training loops (pretraining, finetuning, refitting, and task adaptation), the model can be mutated. Fields can be added (model.extend), removed (model.delete), updated (model.update / with model.override), and reset (model.reset). See the model update guide.
• Production semantics for missingness. null, padded, masked, and valued are distinct states in the tensorfield type system. They are not collapsed into one generic missing-value bucket.
• Online state lives with the model. Stateful components such as category vocabularies, counters, and numeric normalization state are learned during streaming training and serialized with checkpoints, so deployment does not depend on a parallel tokenizer or normalizer artifact.
• Training-serving parity. The same configured graph is used for fitting, validation, testing, batch prediction, and LitServe-backed online inference.
• Target-trained counterfactuals. Training can periodically remove whole field instances with target=True or p_prune, not just mask individual values. At inference time, schema overrides support ablation questions such as "what changes if device data is unavailable?" without retraining a separate model for every feature-removal scenario.

Where It Fits

Use json2vec when the hierarchy is part of the signal:

• customer, account, transaction, statement, device, and session records
• flight itineraries, legs, segments, and events
• orders, shipments, fulfillment events, and support histories
• entities with repeated sub-objects, evolving schemas, and mixed datatypes
• embedding retrieval, anomaly detection, counterfactual ablation, and multi-target prediction over nested records

For more context on the modeling problem, read Why JSON2Vec.

What It Does Not Do

json2vec stops at the representation and typed prediction layer. It does not try to be a feature store, governance system, rule engine, authorization layer, decision-capture system, or audit platform. Those systems can consume json2vec embeddings and predictions, but their policies and operational controls remain separate concerns.

It also does not require users to publish data, schemas, checkpoints, or model parameters. The open-source layer is the reusable encoder and runtime infrastructure. Your data stays yours, and so do your parameters. The framework works under the assumption that model parameters will not be shared.

What Is In This Repository

This repository currently contains:

• the core library under src/json2vec/
• tensorfield plugins for number, category, set, dateparts, entity, vector, and text
• a preprocessor registry for dataset-specific preprocessing
• a LitServe deployment entrypoint for serving from checkpoints
• tests covering structure loading, data processing, tensorfields, training helpers, logging, and inference
• rendered tutorial and guide notebooks under docs/
• diagrams plus whitepaper in docs/

Install

For local development:

uv sync

The package requires Python >=3.12.

Hello World Notebook

The hello world notebook trains a tiny model from the bundled Iris JSONL buffer. It demonstrates the full loop: create a Polars DataFrame, declare a schema, train a supervised category target, then call predict and embed.

import lightning.pytorch as lit
import polars as pl
import torch
from rich.pretty import pprint

import json2vec as j2v


records = pl.read_ndjson("docs/data/iris.jsonl").head(36)

model = j2v.Model.from_schema(
    j2v.Number("sepal_length"),
    j2v.Number("petal_length"),
    j2v.Category("species", target=True, max_vocab_size=4, topk=[2]),
    d_model=16,
    n_layers=1,
    n_heads=4,
    batch_size=8,
    embed=True,
    optimizer=lambda module: torch.optim.AdamW(module.parameters(), lr=1e-2),
)

datamodule = j2v.PolarsDataModule.from_model(
    model,
    train=records,
    validate=records,
    num_workers=0,
    persistent_workers=False,
    pin_memory=False,
    observation_buffer_size=32,
    sample_rate=1.0,
)

trainer = lit.Trainer(
    accelerator="cpu",
    max_epochs=1,
    logger=False,
    enable_progress_bar=False,
    enable_model_summary=False,
    enable_checkpointing=False,
    limit_train_batches=1,
    limit_val_batches=1,
)

trainer.fit(model=model, datamodule=datamodule)

batch = [[record] for record in records.to_dicts()[:3]]

pprint(model.predict(batch))
pprint(model.embed(batch))

The prediction call returns a typed result for record/species. The embedding call returns the configured record embedding for each input observation.

Documentation

The tutorial examples live as self-contained notebooks under docs/ and are rendered in the documentation site. Build the site locally with:

uv run --extra docs mkdocs build --strict

Useful entry points:

• Getting Started
• Why JSON2Vec
• Schemas & Queries
• Model Updates
• Hello World
• Masked Pretraining
• Nested Supervised Training
• Supervised Tabular Training
• Field Ablation
• Preprocessors
• Tensorfield Extensions
• Serving
• API Reference
• Whitepaper

Core Concepts

• Model.from_schema(...) builds the model tree plus masking, targeting, and embedding controls.
• Array nodes describe hierarchical grouping and aggregation.
• Field Request nodes declare a type, a query, and type-specific options.
• Address values are stable paths such as record/account/transaction/amount.
• jmespath queries extract values from each observation.
• TensorField instances preserve typed content plus state tokens such as valued, null, padded, and masked.
• Parcel objects carry embeddings from leaves to parent arrays and then up the tree.
• heritage is the path from a leaf to the root; decoders use that path when reconstructing masked, targeted, or supervised targets.

For large local or cloud-hosted datasets, StreamingDataModule supports these dataset suffixes:

• ndjson
• parquet
• feather
• avro
• csv
• orc
• json

Supported dataset roots are local paths and s3://... URIs.

How The Graph Runs

For each batch:

1. Each field request extracts values with its jmespath query.
2. The matching tensorfield plugin tensorizes values, updates online state when allowed for the current split, and records trainable targets when masking or targeting occurs.
3. Leaf embedders emit parcels to their parent arrays.
4. Array nodes run bottom-up, aggregate child parcels, and emit parent context.
5. Leaf decoders consume their context path to reconstruct trainable targets.

Random p_mask corrupts individual values. Random p_prune removes whole field instances across an observation. target=True is shorthand for p_prune=1.0; embed=True exposes embeddings during prediction.

Preprocessor Model

Preprocessors are optional registered Python callables. See the preprocessor guide for examples. If no preprocessor is configured, each observation is used as-is without calling a default function.

Custom preprocessors are registered with @preprocess(yields=False) for single-object transformations or @preprocess(yields=True) for generators.

• transformation preprocessors must return a single dict
• generator preprocessors may yield dict objects or return a list[dict]
• every emitted object is wrapped as a single-item root array before tensorization

Configured dataset.kwargs are passed into the preprocessor, with unsupported keyword arguments automatically ignored.

Tensorfield Plugins

Each tensorfield plugin provides a request schema plus the model components needed to encode values, decode predictions, compute losses, and optionally serialize outputs. See Tensorfield Extensions for a custom plugin walkthrough. Built-in tensorfields share the base leaf options name, query, pooling, weight, n_heads, n_linear, dropout, p_mask, and p_prune.

Type	Use It For	Key Options
number	Scalar numeric values. Values are padded with explicit state tokens, normalized online during training, embedded with learned Fourier features, and decoded as regression targets.	jitter, n_bands, offset, alpha, objective (mae, mse, huber)
category	Single-label categorical values with an online vocabulary stored in the checkpoint. Unknown or overflow labels route to a reserved unavailable bucket instead of becoming null. Prediction output includes label probabilities and optional top-k candidates.	max_vocab_size, n_bands, p_unavailable, topk
set	Unordered collections of categorical labels, encoded as a multi-hot vector over an online vocabulary. Strings are treated as one-item sets, iterables as many-item sets, and unknown labels use the reserved unavailable bucket.	max_vocab_size, p_unavailable
dateparts	Datetime values represented through selected calendar/time components. Inputs may be native datetimes or strings parsed with a configured pattern.	dateparts (day_of_year, week_of_year, month_of_year, day_of_month, week_of_month, day_of_week, hour_of_day, minute_of_hour), pattern
entity	Hashable identifiers where the useful signal is equality or co-occurrence within the current observation rather than a global vocabulary. Values are re-indexed locally per observation and require at least two slots per observation.	topk
vector	Fixed-width numeric embeddings or dense feature vectors supplied by another model or system. Inputs may be lists, tuples, 1D NumPy arrays, or 1D Torch tensors and are projected into d_model.	n_dim, objective (l1, l2)
text	String values encoded by a frozen Hugging Face AutoModel, pooled, and projected into d_model. Masked or targeted text is trained by reconstructing the encoder representation rather than generating text.	model_name, max_length, encoder_batch_size, encoder_pooling (cls, mean, pooler), objective (l1, l2), revision, local_files_only

The text tensorfield requires the optional transformers dependency and is not installed by default:

uv sync --extra text

Community

Join the Discord channel for questions, design discussion, and release notes: https://discord.gg/DVyZUkvTFA

Repository Layout

• src/json2vec/architecture: model assembly, attention, pooling, and parcel routing
• src/json2vec/data: dataset fetch/read/process/batch/encode pipeline
• src/json2vec/inference: serving and prediction callbacks
• src/json2vec/logging: runtime logging callbacks
• src/json2vec/preprocessors: preprocessor registry
• src/json2vec/structs: pydantic config models, enums, and tree nodes
• src/json2vec/tensorfields: tensorfield plugin system and built-in field types
• tests/: package test suite
• docs/whitepaper.typ: longer written documentation

Development

Run the test suite with:

uv run pytest

Run lint checks with:

uv run ruff check

License

Licensed under the Apache License, Version 2.0. See LICENSE and NOTICE.

References

• BIBLIOGRAPHY.md
• CITATION.bib