Metadata-Version: 2.4
Name: json2vec
Version: 0.4.2
Summary: {...} -> [*]
License-Expression: Apache-2.0
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: beartype>=0.21.0
Requires-Dist: pluggy>=1.6.0
Requires-Dist: rich>=14.0.0
Requires-Dist: pydantic>=2.11.7
Requires-Dist: jmespath>=1.0.1
Requires-Dist: loguru>=0.7.3
Requires-Dist: anytree>=2.13.0
Requires-Dist: ordered-set>=4.1.0
Requires-Dist: pyarrow>=21.0.0
Requires-Dist: polars>=1.35.2
Requires-Dist: numpy>=2.2.6
Requires-Dist: lightning>=2.6.4
Requires-Dist: tensordict>=0.10.0
Requires-Dist: torch>=2.7.1
Provides-Extra: serving
Requires-Dist: litserve>=0.2.13; extra == "serving"
Requires-Dist: pydantic-settings>=2.10.1; extra == "serving"
Provides-Extra: text
Requires-Dist: transformers>=4.55.0; extra == "text"
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.6; extra == "docs"
Requires-Dist: mkdocs-jupyter>=0.26.3; extra == "docs"
Requires-Dist: mkdocstrings[python]>=0.27; extra == "docs"
Dynamic: license-file

<p align="center">
  <img src="https://json2vec.github.io/json2vec/diagrams/json2vec.png" alt="JSON2Vec logo" width="180">
</p>

<h1 align="center">JSON2Vec</h1>

<p align="center">
  <img alt="Python 3.12+" src="https://img.shields.io/badge/python-3.12%2B-3776AB?logo=python&amp;logoColor=white" />
  <a href="LICENSE"><img alt="Apache-2.0 license" src="https://img.shields.io/badge/license-Apache--2.0-2E8B57" /></a>
  <a href="https://json2vec.github.io/json2vec/"><img alt="Documentation" src="https://img.shields.io/badge/docs-MkDocs-526CFE?logo=materialformkdocs&amp;logoColor=white" /></a>
  <!-- discord-invite:start -->
  <a href="https://discord.gg/DVyZUkvTFA"><img alt="Discord channel invite" src="https://img.shields.io/badge/discord-join%20the%20channel-5865F2?logo=discord&amp;logoColor=white" /></a>
  <!-- discord-invite:end -->
</p>

`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](https://json2vec.github.io/json2vec/guides/tensorfields/) 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](https://json2vec.github.io/json2vec/guides/model-update/).
- **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](https://json2vec.github.io/json2vec/motivation/).

## 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/`](https://json2vec.github.io/json2vec/)
- diagrams plus whitepaper in [`docs/`](https://json2vec.github.io/json2vec/)

## Install

For local development:

```bash
uv sync
```

The package requires Python `>=3.12`.

## Hello World Notebook

The [hello world notebook](https://json2vec.github.io/json2vec/tutorials/hello-world/) 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`.

```python
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:

```bash
uv run --extra docs mkdocs build --strict
```

Useful entry points:

- [Getting Started](https://json2vec.github.io/json2vec/getting-started/)
- [Why JSON2Vec](https://json2vec.github.io/json2vec/motivation/)
- [Schemas & Queries](https://json2vec.github.io/json2vec/guides/model-schemas/)
- [Model Updates](https://json2vec.github.io/json2vec/guides/model-update/)
- [Hello World](https://json2vec.github.io/json2vec/tutorials/hello-world/)
- [Masked Pretraining](https://json2vec.github.io/json2vec/tutorials/pretraining/)
- [Nested Supervised Training](https://json2vec.github.io/json2vec/tutorials/nested-supervised-training/)
- [Supervised Tabular Training](https://json2vec.github.io/json2vec/tutorials/supervised-tabular-training/)
- [Field Ablation](https://json2vec.github.io/json2vec/guides/field-ablation/)
- [Preprocessors](https://json2vec.github.io/json2vec/guides/preprocessors/)
- [Tensorfield Extensions](https://json2vec.github.io/json2vec/guides/tensorfields/)
- [Serving](https://json2vec.github.io/json2vec/tutorials/serving/)
- [API Reference](https://json2vec.github.io/json2vec/reference/api/)
- [Whitepaper](https://json2vec.github.io/json2vec/whitepaper.pdf)

## 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](https://json2vec.github.io/json2vec/guides/preprocessors/) 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](https://json2vec.github.io/json2vec/guides/tensorfields/)
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:

```bash
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`](https://json2vec.github.io/json2vec/whitepaper.pdf): longer written documentation

## Development

Run the test suite with:

```bash
uv run pytest
```

Run lint checks with:

```bash
uv run ruff check
```

## License

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

## References

- `BIBLIOGRAPHY.md`
- `CITATION.bib`
