AI Quickstart

This page gives AI agents and experienced Python users the shortest reliable path through json2vec. If you are new to the package, start with Getting Started.

Import Style

Use the package root for normal work:

import json2vec as jv

Build A Model

Top-level records use field names directly. json2vec infers request queries from those names.

model = jv.Model(
    d_model=64,
    n_layers=2,
    n_heads=4,
    batch_size=32,
    amount=jv.Number,
    merchant=jv.Category(size=4096),
    label=jv.Category(target=True, size=2),
)

model

target=True is exact shorthand for p_prune=1.0: it withholds label from the input and trains the decoder to reconstruct it from the remaining fields. It is the “always hidden” supervised-target case; use p_mask for stochastic self-supervised masking with rates lower than 1.0.

Add Nested Data

Use Branch to group fields that share a repeated context, such as line items.

model = jv.Model(
    d_model=64,
    n_layers=2,
    n_heads=4,
    line_items=jv.Branch(
        length=32,
        sku=jv.Category(size=2048),
        quantity=jv.Number,
        price=jv.Number,
    ),
    returned=jv.Category(target=True, size=2),
)

model

For ordered histories, set overflow="tail" on the branch when the newest records should be retained after length is reached. Use overflow="error" when truncation should fail.

This reads records shaped like:

{
    "line_items": [
        {"sku": "ABC-123", "quantity": 2, "price": 19.99},
        {"sku": "XYZ-999", "quantity": 1, "price": 7.50},
    ],
    "returned": "false",
}

Mask Repeated Windows

Use jv.Mask(...) on a Branch when reconstruction should be sampled by repeated coordinate instead of independent leaf value.

model = jv.Model(
    d_model=64,
    n_layers=2,
    n_heads=4,
    embed=True,
    events=jv.Branch(
        length=64,
        overflow="tail",
        mask=jv.Mask(rate=0.25, window=16),
        event_type=jv.Category(size=128),
        amount=jv.Number,
    ),
)

model

Use Dynamic Masking for rate, count, window, offset, and exclude semantics.

Name The Root Context

The generated root branch is named record by default. Override it with name=....

model = jv.Model(
    name="event",
    d_model=32,
    n_layers=1,
    n_heads=4,
    embed=True,
    amount=jv.Number,
)

model

Train Quickly

Use PolarsDataModule(...) for in-memory examples. jv.Model is a Lightning module, so training uses the normal Lightning Trainer.fit(...) loop.

import lightning.pytorch as lit
import polars as pl
import torch

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

model = jv.Model(
    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),
    sepal_length=jv.Number,
    petal_length=jv.Number,
    species=jv.Category(target=True, size=4, topk=[2]),
)

datamodule = jv.PolarsDataModule(
    model=model,
    train=records,
    validate=records,
    num_workers=0,
    persistent_workers=False,
    pin_memory=False,
)

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

For larger file-backed datasets, use Data Modules to choose PolarsDataModule or StreamingDataModule.

Predict

model.predict(...) accepts a list of raw dictionaries. Configured embeddings appear in the same address-keyed output.

predictions = model.predict(records.to_dicts()[:3])

species = predictions[jv.Address("record", "species")]
record = predictions[jv.Address("record")]

For batch prediction jobs, use Batch Inference: configure a predict split on a data module, attach jv.Writer(...) to Trainer(callbacks=[...]), and call trainer.predict(...).

Shape Raw Data

Use query=... when the source shape is stable but does not match the schema name.

model = jv.Model(
    d_model=32,
    n_layers=1,
    n_heads=4,
    amount=jv.Number(query='[*].transaction."amount_usd"'),
    merchant=jv.Category(query="[*].transaction.merchant_name", size=1024),
    label=jv.Category(query="[*].outcome", target=True, size=2),
)

model

Use a preprocessor when the source needs Python logic before querying. The decorator returns a Preprocessor object that can be passed to data modules, deployment configuration, model.encode(...), or model.predict(...).

@jv.preprocess
def normalize_request(record: dict) -> jv.Observation:
    return jv.Observation({
        "amount": float(record["amount_usd"]),
        "merchant": record["merchant"].strip().lower(),
        "label": record["label"],
    })

Then pass it to data modules, model.encode(...), model.predict(...), or deployments where supported. See Data Modules for preprocessor execution in train, validate, test, and predict splits.

Reshape Outputs

Use a postprocessor to turn address-keyed predictions into an API or warehouse shape.

@jv.postprocess
def risk_response(predictions):
    label = {}
    for address, values in predictions.items():
        if str(address) == "record/label":
            label = values.get("content", {})
            break

    return {
        "risk": {
            "prediction": label.get("value"),
            "scores": label.get("probability"),
        }
    }

Serve

Serving helpers are top-level exports when the serving extra is installed.

deployment = (
    jv.Deployment(model=model, accelerator="cpu", max_batch_size=16)
    .postprocess(risk_response)
    .update(jv.where("name") == "label", topk=[2])
)

Use model=... for an already constructed model during local setup, or checkpoint="model.ckpt" when loading a trained checkpoint. Call deployment.serve() from an application entry point.

Rules For Agents

• Prefer import json2vec as jv.
• Prefer Model(...) and top-level exports.
• Use name=... to name the generated root context.
• Set size on categorical fields where necessary.
• Only put Entity under repeated contexts.
• Use embed=True for embeddings and target=True for supervised targets.
• Use p_mask for stochastic self-supervised reconstruction; use target=True or p_prune=1.0 when a field should always be hidden and decoded.
• Use jv.Mask(...) on branches for recency-window or row-level reconstruction.
• Use Lightning Trainer with PolarsDataModule or StreamingDataModule for training and batch prediction.
• Use jv.Writer for Parquet batch inference output.

Use Getting Started and the guide pages for runnable patterns.