Metadata-Version: 2.4
Name: torchaid
Version: 0.1.0
Summary: A PyTorch utility library that streamlines the deep learning training pipeline.
Author: Harunori Kawano
License: MIT
Project-URL: Homepage, https://github.com/HarunoriKawano/torchaid
Project-URL: Issues, https://github.com/HarunoriKawano/torchaid/issues
Keywords: pytorch,deep learning,training,transformer
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: tqdm>=4.60.0
Requires-Dist: pandas>=1.5.0
Requires-Dist: scikit-learn>=1.0.0
Requires-Dist: numpy>=1.21.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Dynamic: license-file

# torchaid

**torchaid** is a PyTorch utility library that provides structured abstractions and reusable components to streamline the deep learning training pipeline.

## Features

- **Structured training abstractions** — type-safe base classes for inputs, outputs, metrics, and settings built on Pydantic v2
- **Training framework** — a full training loop with mixed-precision support, automatic checkpointing, metric logging (CSV), and early stopping
- **Transformer modules** — standard and relative-position-aware Transformer encoder layers with multi-head self-attention
- **Task templates** — ready-to-use implementation for multi-label classification
- **Utilities** — dataset splitting, random seed management, attention mask generation, and JSON-to-Pydantic loading
- **Learning rate schedulers** — cosine decay with linear warm-up, and triangular2 cyclic scheduling

## Requirements

- Python 3.10+
- PyTorch 2.0+

## Installation

```bash
pip install torchaid
```

Or install from source:

```bash
git clone https://github.com/harunori-kawano/torchaid.git
cd torchaid
pip install -e .
```

## Quick Start

### 1. Define your data schema

```python
from torchaid import BaseInputs, BaseOutputs
import torch

class MyInputs(BaseInputs):
    input_ids: torch.Tensor   # (B, L)
    labels: torch.LongTensor  # (B,)

class MyOutputs(BaseOutputs):
    logits: torch.Tensor      # (B, num_classes)
```

### 2. Implement your model

```python
from torchaid import TaskModule, Mode, BaseOutputs
from torch import nn

class MyModel(TaskModule):
    def __init__(self, vocab_size: int, num_classes: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, 128)
        self.classifier = nn.Linear(128, num_classes)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, mode: Mode, batch: MyInputs) -> BaseOutputs:
        x = self.embed(batch.input_ids).mean(dim=1)
        logits = self.classifier(x)
        loss = self.criterion(logits, batch.labels)
        return MyOutputs(loss=loss, logits=logits)
```

### 3. Define metrics and settings

```python
from torchaid import BaseMetrics, BaseSettings, BaseMetricCalculator
from typing import Optional, Any

class MyMetrics(BaseMetrics):
    train_loss: Optional[float] = None
    val_loss: Optional[float] = None

class MySettings(BaseSettings):
    batch_size: int = 32
    max_epoch_num: int = 10

class MyCalculator(BaseMetricCalculator[MyMetrics]):
    def __init__(self):
        super().__init__(MyMetrics())
        self._losses: list[float] = []

    def train_step(self, outputs, batch) -> dict[str, Any]:
        loss = outputs.loss.item()
        self._losses.append(loss)
        return {"loss": loss}

    def val_step(self, outputs, batch) -> dict[str, Any]:
        return self.train_step(outputs, batch)

    def test_step(self, outputs, batch) -> dict[str, Any]:
        return self.train_step(outputs, batch)

    def check(self) -> bool:
        import statistics
        self.metrics.train_loss = statistics.mean(self._losses)
        return True

    def test(self): pass

    def reset(self):
        self._losses.clear()
```

### 4. Train

```python
import torch
from torchaid.core.trainer import TrainFramework

settings = MySettings(batch_size=32, max_epoch_num=10, device="cuda")
model = MyModel(vocab_size=1000, num_classes=5)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

framework = TrainFramework(
    model=model,
    ls=settings,
    metric_calculator=MyCalculator(),
    optimizer=optimizer,
    inputs_config=MyInputs,
)

framework.train(train_dataset, val_dataset, save_dir="./outputs")
```

## Module Overview

| Module | Description |
|--------|-------------|
| `torchaid.core` | Base classes (`BaseInputs`, `BaseOutputs`, `BaseMetrics`, `BaseSettings`, `TaskModule`, `BaseMetricCalculator`, `Mode`) and `TrainFramework` |
| `torchaid.templates.multilabel_classification` | Complete template for multi-label classification |
| `torchaid.extras.modules.transformer` | `Transformer`, `TransformerWithRelativePosition`, and sub-modules |
| `torchaid.extras.modules.positional_encoders` | `PositionalEmbedding`, `RelativePositionEmbedding` |
| `torchaid.extras.utils` | `split_dataset`, `set_random_seed`, `make_attention_mask`, `json_to_instance` |
| `torchaid.extras.scheduler` | `get_cosine_scheduler`, `get_cycle_scheduler` |

## Template: Multi-Label Classification

```python
from torchaid.templates import multilabel_classification as mlc
from torchaid.core.trainer import TrainFramework
from torch import nn
import torch

backbone = nn.Sequential(nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10))
model = mlc.MultiLabelClassification(backbone)
optimizer = torch.optim.Adam(model.parameters())

framework = TrainFramework(
    model=model,
    ls=settings,
    metric_calculator=mlc.MetricsCalculator(),
    optimizer=optimizer,
    inputs_config=mlc.Inputs,
)
```

## Extras

### Cosine Decay Scheduler

```python
from torchaid.extras.scheduler import get_cosine_scheduler

scheduler = get_cosine_scheduler(
    optimizer, warmup_steps=500, max_steps=10000
)
```

### Dataset Split

```python
from torchaid.extras.utils import split_dataset

train, val, test = split_dataset(dataset, ratios=[8, 1, 1], seed=42)
```

### Relative Position Transformer

```python
from torchaid.extras.modules.transformer import TransformerWithRelativePosition

layer = TransformerWithRelativePosition(
    hidden_size=256,
    intermediate_size=1024,
    num_attention_heads=8,
    dropout_probability=0.1,
    max_length=512,
    with_cls=True,
)
```

## License

MIT License. See [LICENSE](LICENSE) for details.
