Batch Inference

model.predict([...]) is the interactive path for a list of raw dictionaries. For large prediction jobs, use a data module with Lightning trainer.predict(...). Attach jv.Writer(...) to write each prediction batch to Parquet as it finishes.

Minimal Batch Prediction

import lightning.pytorch as lit
import polars as pl

import json2vec as jv

predict_frame = pl.DataFrame(
    {
        "amount": [12.5, 8.0],
        "merchant": ["books", "coffee"],
    }
)

datamodule = jv.PolarsDataModule(
    model=model,
    predict=predict_frame,
    num_workers=0,
    persistent_workers=False,
    pin_memory=False,
)

writer = jv.Writer("predictions")

trainer = lit.Trainer(
    accelerator="cpu",
    callbacks=[writer],
    logger=False,
)

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

Writer is a Lightning prediction callback. It belongs in Trainer(callbacks=[...]) and runs during trainer.predict(...).

Output Files

Writer creates one Parquet file per local rank:

predictions/
  rank-0.parquet
  rank-1.parquet

Single-process jobs usually write only rank-0.parquet. Distributed jobs write one file per local rank, and downstream jobs should read all rank files.

Each row contains:

  • inputs: row metadata captured during encoding.
  • predictions: a nested struct keyed by schema address.

Target predictions are included for configured targets. Embeddings are included for nodes configured with embed=True.

Writer Options

Argument Meaning
path Local or mounted output directory for rank-partitioned Parquet files.
flush_every_n_batches Optional Parquet writer flush interval.
postprocessor Optional callable that rewrites prediction dictionaries before write. 
writer = jv.Writer(
    "predictions",
    flush_every_n_batches=100,
)

Postprocess Before Writing

Use a postprocessor when downstream consumers need flat columns, renamed addresses, filtered payloads, or fewer fields.

import json2vec as jv


@jv.postprocess
def compact_rows(predictions, *, metadata):
    label = predictions.get(jv.Address("record", "label"), {}).get("content", {})

    return {
        jv.Address("record", "warehouse_row"): {
            "request_id": [item.get("request_id") for item in metadata],
            "label": label.get("value"),
            "label_probability": label.get("probability"),
        }
    }


writer = jv.Writer(
    "warehouse-predictions",
    postprocessor=compact_rows,
)

Keep returned lists row-aligned with the prediction batch. That lets the writer build a valid Parquet table.

Writer postprocessors can request named runtime values such as:

Parameter Meaning
input Encoded Lightning batch.
batch Alias for the encoded batch.
metadata Raw input metadata for each row.
batch_indices Lightning batch indices when provided.
batch_idx Prediction batch index.
dataloader_idx Prediction dataloader index.

For more output-shaping patterns, see Postprocessors.

Streaming Prediction

For file-backed local or S3 inputs, use StreamingDataModule with a predict split:

import lightning.pytorch as lit

import json2vec as jv

datamodule = jv.StreamingDataModule(
    model=model,
    root="s3://my-bucket/events",
    suffix="parquet",
    predict=r"/predict/.*\.parquet$",
    sharding="chunk",
)

trainer = lit.Trainer(
    accelerator="gpu",
    devices=1,
    callbacks=[jv.Writer("predictions")],
)

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

StreamingDataModule can read from S3. Writer writes to a local or mounted filesystem path, so copy or sync those rank files downstream when the final destination is object storage.

Choosing An Inference Path

Need Use
Debug one request shape model.predict([...])
Notebook demo model.predict([...])
Batch DataFrame prediction PolarsDataModule(..., predict=...) plus Writer
Batch file or S3 prediction StreamingDataModule(..., predict=...) plus Writer
Online service jv.Deployment

Common Mistakes

  • Forgetting to configure a predict split on the data module.
  • Calling model.predict(...) in a Python loop for a large batch job.
  • Returning non-row-aligned lists from a writer postprocessor.
  • Reading only one rank-*.parquet file from a distributed prediction job.
  • Using a different preprocessor at prediction time than the one used during training.

Where Next