coffeetrain 0.1.0__tar.gz

coffeetrain-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,37 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+      - master
+
+jobs:
+  lint-and-test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install package and dev dependencies
+        run: uv sync --group dev
+
+      - name: Lint (ruff)
+        run: uv run ruff check .
+
+      - name: Test (pytest)
+        run: uv run pytest -q

coffeetrain-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,65 @@
+name: Publish
+
+on:
+  push:
+    tags:
+      - "v*"
+
+jobs:
+  pypi-publish:
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+      contents: read
+    environment:
+      name: pypi
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Setup uv
+        uses: astral-sh/setup-uv@v6
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Verify tag matches project version
+        run: |
+          python - <<'PY'
+          import os
+          import pathlib
+          import tomllib
+
+          ref_name = os.environ["GITHUB_REF_NAME"]
+          if not ref_name.startswith("v"):
+              raise SystemExit(f"Expected tag starting with 'v', got: {ref_name}")
+
+          tag_version = ref_name[1:]
+          pyproject_path = pathlib.Path("pyproject.toml")
+          data = tomllib.loads(pyproject_path.read_text(encoding="utf-8"))
+          project_version = data["project"]["version"]
+
+          if tag_version != project_version:
+              raise SystemExit(
+                  f"Tag version '{tag_version}' does not match pyproject version '{project_version}'."
+              )
+
+          print(f"Tag version matches pyproject version: {project_version}")
+          PY
+
+      - name: Install build tooling
+        run: uv pip install --system build twine
+
+      - name: Build distributions
+        run: python -m build
+
+      - name: Verify distributions
+        run: twine check dist/*
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          verbose: true

coffeetrain-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

coffeetrain-0.1.0/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.4
+Name: coffeetrain
+Version: 0.1.0
+Summary: Lightweight event-driven PyTorch trainer with composable callbacks
+Project-URL: Repository, https://github.com/paul-english/coffeetrain
+Author: Paul English
+License: Apache-2.0
+Keywords: callbacks,ml,pytorch,training
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.12
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: torch
+Requires-Dist: torchmetrics>=1.8.2
+Requires-Dist: tqdm
+Provides-Extra: comet
+Requires-Dist: comet-ml; extra == 'comet'
+Provides-Extra: optimi
+Requires-Dist: torch-optimi; extra == 'optimi'
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.18.0; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+# coffeetrain
+
+Lightweight event-driven PyTorch trainer with composable callbacks. Inspired by [MosaicML Composer](https://github.com/mosaicml/composer) but implemented fewer external dependencies and allowing a newer PyTorch version (2.10 as of this point).
+
+## Features
+
+- **Event-driven lifecycle**: `fit_start`, `epoch_start`, `batch_start`, `before_forward`, `after_forward`, `before_loss`, `after_loss`, `before_backward`, `after_backward`, `batch_end`, `eval_*`, etc.
+- **Composable callbacks**: EMA, SWA, checkpointing, W&B, Comet, early stopping, batch size scheduling, and more.
+- **TrainerModel protocol**: Simple interface (`forward`, `loss`) for model integration.
+- **Accelerate support**: Distributed training via HuggingFace Accelerator.
+
+## Installation
+
+```bash
+pip install coffeetrain
+```
+
+Optional extras:
+```bash
+pip install coffeetrain[wandb,comet,optimi]
+```
+
+## Quick Start
+
+```python
+from coffeetrain import Trainer, CosineWarmupScheduler, HistoryCallback, BestModelCheckpointer
+from coffeetrain import create_optimizer
+from coffeetrain.optimizers import OptimizerConfig
+
+model = MyModel()
+optimizer = create_optimizer(model.parameters(), OptimizerConfig(name="adamw", lr=1e-4, weight_decay=0.01))
+scheduler = CosineWarmupScheduler(optimizer, warmup_steps=100, total_steps=1000)
+
+trainer = Trainer(
+    model=model,
+    train_dataloader=train_loader,
+    optimizers=optimizer,
+    schedulers=scheduler,
+    max_epochs=10,
+    callbacks=[
+        HistoryCallback(save_dir="output"),
+        BestModelCheckpointer(save_dir="output", metric_name="loss", mode="min"),
+    ],
+)
+trainer.fit()
+```
+
+## Callbacks
+
+| Callback | Description |
+|----------|-------------|
+| `BestModelCheckpointer` | Save best model by metric |
+| `HistoryCallback` | Track and save training history to JSON |
+| `EMACallback` | Exponential moving average of weights |
+| `SWACallback` | Stochastic weight averaging |
+| `EarlyStoppingCallback` | Stop when metric stops improving |
+| `WandbCallback` | Log to Weights & Biases |
+| `CometCallback` | Log to Comet.ml |
+| `BatchSizeSchedulerCallback` | Batch size warmup |
+| `ScheduleLoggerCallback` | Log LR schedule phase transitions |
+| `ParameterCounter` | Print parameter counts at start |
+| `SpeedMonitor` | Track samples/sec |
+| `ProgressCallback` | Print epoch summaries |
+| `LRMonitor` | Log learning rates |
+| `TorchMetricsCallback` | Integrate torchmetrics |
+
+## License
+
+Apache-2.0
+
+## Tests
+
+From repository root:
+
+```bash
+uv run pytest packages/coffeetrain/tests -q
+```

coffeetrain-0.1.0/README.md

@@ -0,0 +1,77 @@
+# coffeetrain
+
+Lightweight event-driven PyTorch trainer with composable callbacks. Inspired by [MosaicML Composer](https://github.com/mosaicml/composer) but implemented fewer external dependencies and allowing a newer PyTorch version (2.10 as of this point).
+
+## Features
+
+- **Event-driven lifecycle**: `fit_start`, `epoch_start`, `batch_start`, `before_forward`, `after_forward`, `before_loss`, `after_loss`, `before_backward`, `after_backward`, `batch_end`, `eval_*`, etc.
+- **Composable callbacks**: EMA, SWA, checkpointing, W&B, Comet, early stopping, batch size scheduling, and more.
+- **TrainerModel protocol**: Simple interface (`forward`, `loss`) for model integration.
+- **Accelerate support**: Distributed training via HuggingFace Accelerator.
+
+## Installation
+
+```bash
+pip install coffeetrain
+```
+
+Optional extras:
+```bash
+pip install coffeetrain[wandb,comet,optimi]
+```
+
+## Quick Start
+
+```python
+from coffeetrain import Trainer, CosineWarmupScheduler, HistoryCallback, BestModelCheckpointer
+from coffeetrain import create_optimizer
+from coffeetrain.optimizers import OptimizerConfig
+
+model = MyModel()
+optimizer = create_optimizer(model.parameters(), OptimizerConfig(name="adamw", lr=1e-4, weight_decay=0.01))
+scheduler = CosineWarmupScheduler(optimizer, warmup_steps=100, total_steps=1000)
+
+trainer = Trainer(
+    model=model,
+    train_dataloader=train_loader,
+    optimizers=optimizer,
+    schedulers=scheduler,
+    max_epochs=10,
+    callbacks=[
+        HistoryCallback(save_dir="output"),
+        BestModelCheckpointer(save_dir="output", metric_name="loss", mode="min"),
+    ],
+)
+trainer.fit()
+```
+
+## Callbacks
+
+| Callback | Description |
+|----------|-------------|
+| `BestModelCheckpointer` | Save best model by metric |
+| `HistoryCallback` | Track and save training history to JSON |
+| `EMACallback` | Exponential moving average of weights |
+| `SWACallback` | Stochastic weight averaging |
+| `EarlyStoppingCallback` | Stop when metric stops improving |
+| `WandbCallback` | Log to Weights & Biases |
+| `CometCallback` | Log to Comet.ml |
+| `BatchSizeSchedulerCallback` | Batch size warmup |
+| `ScheduleLoggerCallback` | Log LR schedule phase transitions |
+| `ParameterCounter` | Print parameter counts at start |
+| `SpeedMonitor` | Track samples/sec |
+| `ProgressCallback` | Print epoch summaries |
+| `LRMonitor` | Log learning rates |
+| `TorchMetricsCallback` | Integrate torchmetrics |
+
+## License
+
+Apache-2.0
+
+## Tests
+
+From repository root:
+
+```bash
+uv run pytest packages/coffeetrain/tests -q
+```

coffeetrain-0.1.0/pyproject.toml

@@ -0,0 +1,59 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "coffeetrain"
+version = "0.1.0"
+description = "Lightweight event-driven PyTorch trainer with composable callbacks"
+readme = "README.md"
+requires-python = ">=3.12"
+authors = [{name = "Paul English"}]
+license = {text = "Apache-2.0"}
+keywords = ["pytorch", "training", "callbacks", "ml"]
+classifiers = [
+  "Development Status :: 4 - Beta",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: Apache Software License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+  "torch",
+  "tqdm",
+  "torchmetrics>=1.8.2",
+  "pydantic>=2.0.0",
+]
+
+[project.optional-dependencies]
+wandb = ["wandb>=0.18.0"]
+comet = ["comet-ml"]
+optimi = ["torch-optimi"]
+
+[project.urls]
+Repository = "https://github.com/paul-english/coffeetrain"
+
+[dependency-groups]
+dev = [
+  "ruff>=0.13.0",
+  "pytest>=9.0.2",
+  "pytest-timeout>=2.4.0",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/coffeetrain"]
+
+[tool.ruff]
+target-version = "py312"
+
+[tool.ruff.lint]
+select = ["E9", "F63", "F7", "F82"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-q"

coffeetrain-0.1.0/src/coffeetrain/__init__.py

@@ -0,0 +1,71 @@
+"""Lightweight event-driven PyTorch trainer with composable callbacks."""
+
+from coffeetrain.events import Event
+from coffeetrain.state import State
+from coffeetrain.callback import Callback
+from coffeetrain.model import TrainerModel, TrainerModelMixin, ModelWrapper
+from coffeetrain.trainer import Trainer
+from coffeetrain.schedulers import CosineWarmupScheduler, LinearWarmupScheduler, WSDScheduler
+from coffeetrain.optimizers import (
+    OptimizerConfig,
+    create_optimizer,
+    create_optimizer_with_param_groups,
+)
+
+# Callbacks
+from coffeetrain.callbacks import (
+    CometCallback,
+    EMACallback,
+    HistoryCallback,
+    BestModelCheckpointer,
+    LRMonitor,
+    ParameterCounter,
+    SpeedMonitor,
+    ProgressCallback,
+    TorchMetricsCallback,
+    WandbCallback,
+    BatchSizeSchedulerCallback,
+    SWACallback,
+    ScheduleLoggerCallback,
+    EarlyStoppingCallback,
+)
+
+# Metrics
+from coffeetrain.metrics import MeanLoss, Perplexity, SafeMulticlassAccuracy
+
+__all__ = [
+    # Core
+    "Event",
+    "State",
+    "Callback",
+    "TrainerModel",
+    "TrainerModelMixin",
+    "ModelWrapper",
+    "Trainer",
+    "CosineWarmupScheduler",
+    "LinearWarmupScheduler",
+    "WSDScheduler",
+    # Optimizers
+    "OptimizerConfig",
+    "create_optimizer",
+    "create_optimizer_with_param_groups",
+    # Callbacks
+    "CometCallback",
+    "EMACallback",
+    "HistoryCallback",
+    "BestModelCheckpointer",
+    "LRMonitor",
+    "ParameterCounter",
+    "SpeedMonitor",
+    "ProgressCallback",
+    "TorchMetricsCallback",
+    "WandbCallback",
+    "BatchSizeSchedulerCallback",
+    "SWACallback",
+    "ScheduleLoggerCallback",
+    "EarlyStoppingCallback",
+    # Metrics
+    "MeanLoss",
+    "Perplexity",
+    "SafeMulticlassAccuracy",
+]

coffeetrain-0.1.0/src/coffeetrain/callback.py

@@ -0,0 +1,120 @@
+"""Callback base class for training hooks."""
+
+from typing import Any, Dict
+
+from coffeetrain.state import State
+
+
+class Callback:
+    """Base class for training callbacks.
+
+    Callbacks receive State at each event and can read/modify it.
+    Override the methods corresponding to events you want to handle.
+
+    All methods are no-ops by default, so subclasses only need to
+    implement the hooks they care about.
+
+    Example:
+        class PrintLossCallback(Callback):
+            def batch_end(self, state: State) -> None:
+                print(f"Loss: {state.loss.item():.4f}")
+    """
+
+    # Initialization
+    def init(self, state: State) -> None:
+        """Called after Trainer construction, before fit().
+
+        Use for callback initialization that needs access to state.
+        """
+        pass
+
+    # Fit lifecycle
+    def fit_start(self, state: State) -> None:
+        """Called at the start of fit(), before any training."""
+        pass
+
+    def fit_end(self, state: State) -> None:
+        """Called at the end of fit(), after all training."""
+        pass
+
+    # Epoch lifecycle
+    def epoch_start(self, state: State) -> None:
+        """Called at the start of each epoch."""
+        pass
+
+    def epoch_end(self, state: State) -> None:
+        """Called at the end of each epoch, before evaluation."""
+        pass
+
+    # Training batch lifecycle
+    def batch_start(self, state: State) -> None:
+        """Called at the start of each training batch."""
+        pass
+
+    def before_forward(self, state: State) -> None:
+        """Called before model.forward()."""
+        pass
+
+    def after_forward(self, state: State) -> None:
+        """Called after model.forward(), state.outputs is set."""
+        pass
+
+    def before_loss(self, state: State) -> None:
+        """Called before model.loss()."""
+        pass
+
+    def after_loss(self, state: State) -> None:
+        """Called after model.loss(), state.loss is set."""
+        pass
+
+    def before_backward(self, state: State) -> None:
+        """Called before loss.backward()."""
+        pass
+
+    def after_backward(self, state: State) -> None:
+        """Called after loss.backward(), before optimizer step."""
+        pass
+
+    def batch_end(self, state: State) -> None:
+        """Called after optimizer step and scheduler step."""
+        pass
+
+    # Evaluation lifecycle
+    def eval_start(self, state: State) -> None:
+        """Called at the start of evaluation."""
+        pass
+
+    def eval_batch_start(self, state: State) -> None:
+        """Called at the start of each eval batch."""
+        pass
+
+    def eval_before_forward(self, state: State) -> None:
+        """Called before model.forward() during eval."""
+        pass
+
+    def eval_after_forward(self, state: State) -> None:
+        """Called after model.forward() during eval."""
+        pass
+
+    def eval_batch_end(self, state: State) -> None:
+        """Called at the end of each eval batch."""
+        pass
+
+    def eval_end(self, state: State) -> None:
+        """Called at the end of evaluation."""
+        pass
+
+    # Checkpointing support
+    def state_dict(self) -> Dict[str, Any]:
+        """Return callback state for checkpointing.
+
+        Override to save callback-specific state.
+        """
+        return {}
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        """Load callback state from checkpoint.
+
+        Override to restore callback-specific state.
+        """
+        pass

coffeetrain-0.1.0/src/coffeetrain/callbacks/__init__.py

@@ -0,0 +1,33 @@
+"""Generic callback implementations for the trainer."""
+
+from coffeetrain.callbacks.batch_size_scheduler import BatchSizeSchedulerCallback
+from coffeetrain.callbacks.checkpointing import BestModelCheckpointer
+from coffeetrain.callbacks.comet import CometCallback
+from coffeetrain.callbacks.early_stopping import EarlyStoppingCallback
+from coffeetrain.callbacks.ema import EMACallback
+from coffeetrain.callbacks.swa import SWACallback
+from coffeetrain.callbacks.history import HistoryCallback
+from coffeetrain.callbacks.lr_monitor import LRMonitor
+from coffeetrain.callbacks.parameter_counter import ParameterCounter
+from coffeetrain.callbacks.progress import ProgressCallback
+from coffeetrain.callbacks.speed_monitor import SpeedMonitor
+from coffeetrain.callbacks.torchmetrics_callback import TorchMetricsCallback
+from coffeetrain.callbacks.wandb import WandbCallback
+from coffeetrain.callbacks.schedule_logger import ScheduleLoggerCallback
+
+__all__ = [
+    "BatchSizeSchedulerCallback",
+    "BestModelCheckpointer",
+    "CometCallback",
+    "EarlyStoppingCallback",
+    "EMACallback",
+    "SWACallback",
+    "HistoryCallback",
+    "LRMonitor",
+    "ParameterCounter",
+    "ProgressCallback",
+    "ScheduleLoggerCallback",
+    "SpeedMonitor",
+    "TorchMetricsCallback",
+    "WandbCallback",
+]