Metadata-Version: 2.4
Name: torch-vault
Version: 0.1.0
Classifier: Programming Language :: Rust
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Dist: torch>=2.0
Requires-Dist: pyarrow>=14.0
Requires-Dist: maturin ; extra == 'dev'
Requires-Dist: pytest ; extra == 'dev'
Requires-Dist: numpy ; extra == 'dev'
Provides-Extra: dev
Summary: Zero-copy PyTorch checkpoint management and SQL-driven dataset streaming
Keywords: pytorch,rust,machine-learning,checkpointing,duckdb
Author-email: Ethan Wu <ewu@u.northwestern.edu>
License: MIT
Requires-Python: >=3.10
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Documentation, https://ethqnol.github.io/torch-vault/
Project-URL: Homepage, https://github.com/ethqnol/torch-vault
Project-URL: Repository, https://github.com/ethqnol/torch-vault

# torch-vault

Zero-copy PyTorch checkpoint management and SQL-driven dataset streaming built with Rust, PyO3, SQLite, and DuckDB.

## Features

- **Safe Checkpointing**: Replaces pickle-based `torch.save` with `safetensors` memory-mapping backed by an atomic SQLite WAL index.
- **Selective Weight Saving**: `vault.save(model, metrics=..., save_if_better="val_loss")` continuously records metric history into SQLite while writing `.safetensors` weight files to disk only when metrics improve.
- **SQL Dataset Streaming**: `Segment("data.parquet").where("quality > 0.8")` runs DuckDB predicate pushdown on disk before streaming records into PyTorch DataLoaders.
- **Zero-Copy Memory Interchange**: Uses the Arrow C Stream Interface (`__arrow_c_stream__` PyCapsule) to pass record batches directly into memory.
- **CLI & WebApp Inspector**: Includes `torch-vault list` and `torch-vault ui` for inspecting checkpoints and charting metrics in a local browser interface.

## Installation

```bash
pip install torch-vault
```

For development build:

```bash
git clone https://github.com/ethqnol/torch-vault.git
cd torch-vault
pip install maturin
maturin develop --release
```

## Quickstart

### Checkpoint Management (`Vault`)

```python
import torch
import torch.nn as nn
from torch_vault import Vault

model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
vault = Vault("./checkpoints")

for epoch in range(1, 11):
    train_loss = 0.5 # compute training loss
    val_loss = 0.4   # compute validation loss

    # Log metrics to SQLite; save weight file on disk only if val_loss improves
    vault.save(
        model=model,
        optimizer=optimizer,
        metrics={"train_loss": train_loss, "val_loss": val_loss},
        epoch=epoch,
        save_if_better="val_loss",
        mode="min",
    )

# Load best model weights automatically by metric
best_weights = vault.load(metric="val_loss", mode="min")
model.load_state_dict(best_weights["model"])
```

### SQL Dataset Streaming & Train/Val/Test Split (`Segment`)

```python
from torch_vault import Segment
from torch_vault.transforms import Compose, SelectKeys, RenameKey, StandardScaler

pipeline = Compose([
    RenameKey(src="score", dst="quality"),
    StandardScaler(keys=["quality"]),
    SelectKeys("quality", "label"),
])

# Stream filtered Parquet data with a 80%/10%/10% split
train_loader, val_loader, test_loader = (
    Segment("data/dataset.parquet")
    .where("quality > 0.7 AND split = 'train'")
    .transform(pipeline)
    .to_dataloader(batch_size=32, split=(0.8, 0.1, 0.1), seed=42)
)
```

### CLI & WebApp Inspector

```bash
# Print terminal ASCII table of logged checkpoints
torch-vault list ./checkpoints

# Launch browser WebApp inspector at http://127.0.0.1:8000/
torch-vault ui ./checkpoints --port 8000
```

## Documentation

Full interactive API reference and architecture specifications are available at [ethqnol.github.io/torch-vault](https://ethqnol.github.io/torch-vault/).

## License

MIT License.

