mltrainer 0.3.0__tar.gz → 0.3.1__tar.gz

mltrainer-0.3.1/.claude/settings.local.json

@@ -0,0 +1,7 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(uv run *)"
+    ]
+  }
+}

mltrainer-0.3.1/.remember/.gitignore

@@ -0,0 +1 @@
+*

mltrainer-0.3.1/.remember/tmp/last-ndc.ts

@@ -0,0 +1 @@
+1784568708

mltrainer-0.3.1/.remember/tmp/last-save-ts

@@ -0,0 +1 @@
+1784568694

mltrainer-0.3.1/.remember/tmp/last-save.json

@@ -0,0 +1 @@
+{"session": "284374ad-dc46-496d-a542-5d141c54654d", "line": 204}

mltrainer-0.3.1/.remember/tmp/save-session.pid

@@ -0,0 +1 @@
+85305

mltrainer-0.3.1/.remember/today-2026-07-20.md

@@ -0,0 +1,2 @@
+## 19:31 | main
+Refactored mltrainer.report() w/ Reporter ABC (RayReporter/MLflow/TensorBoard); type hints→builtin gen; added _to_device helper; fixed ty checks (transformer.py/trainer.py); Makefile: ruff/ty; cleaned ~115 commented LOC (settings.py); 47 tests pass; ready v0.3.1

mltrainer-0.3.1/Makefile

@@ -0,0 +1,12 @@
+format:
+	ruff format src
+	ruff check --fix src
+
+lint:
+	ruff check src
+	ty check src
+
+test:
+	pytest
+
+.PHONY: format lint test

{mltrainer-0.3.0 → mltrainer-0.3.1}/PKG-INFO

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: mltrainer
-Version: 0.3.0
+Version: 0.3.1
 Summary: toolkit for training pytorch models
+Project-URL: GitHub, https://github.com/raoulg/mltrainer
 Author-email: "R.Grouls" <Raoul.Grouls@han.nl>
 License: MIT
 Classifier: Development Status :: 4 - Beta

{mltrainer-0.3.0 → mltrainer-0.3.1}/pyproject.toml

@@ -1,7 +1,7 @@
 
 [project]
 name = "mltrainer"
-version = "0.3.0"
+version = "0.3.1"
 description = "toolkit for training pytorch models"
 
 authors = [
@@ -44,7 +44,7 @@ build-backend = "hatchling.build"
 [tool.hatch.build.targets.wheel]
 packages = ["src/mltrainer"]
 
-[tools.urls]
+[project.urls]
 GitHub = "https://github.com/raoulg/mltrainer"
 
 [dependency-groups]

{mltrainer-0.3.0 → mltrainer-0.3.1}/src/mltrainer/__init__.py

@@ -3,4 +3,4 @@ from mltrainer.trainer import Trainer
 
 __all__ = ["Trainer", "TrainerSettings", "ReportTypes"]
 
-__version__ = "0.2.7"
+__version__ = "0.3.1"

{mltrainer-0.3.0 → mltrainer-0.3.1}/src/mltrainer/metrics.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Iterator, Union
+from typing import Iterator, Literal, Union
 
 import numpy as np
 import torch
@@ -86,7 +86,12 @@ class Accuracy(Metric):
 
 
 class F1Score(Metric):
-    def __init__(self, average="micro", threshold=0.5, epsilon=1e-7):
+    def __init__(
+        self,
+        average: Literal["micro", "macro"] = "micro",
+        threshold: float = 0.5,
+        epsilon: float = 1e-7,
+    ) -> None:
         self.average = average
         self.threshold = threshold
         self.epsilon = epsilon

mltrainer-0.3.1/src/mltrainer/settings.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+from loguru import logger
+from pydantic import BaseModel, ConfigDict, field_validator
+
+
+class FileTypes(Enum):
+    JPG = ".jpg"
+    PNG = ".png"
+    TXT = ".txt"
+    ZIP = ".zip"
+    TGZ = ".tgz"
+    TAR = ".tar.gz"
+    GZ = ".gz"
+
+
+class ReportTypes(Enum):
+    TOML = "TOML"
+    TENSORBOARD = "TENSORBOARD"
+    MLFLOW = "MLFLOW"
+    RAY = "RAY"
+
+
+class FormattedBase(BaseModel):
+    def __str__(self) -> str:
+        return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
+
+    def __repr__(self) -> str:
+        return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
+
+
+class TrainerSettings(FormattedBase):
+    epochs: int
+    metrics: List[Callable]
+    logdir: Path
+    train_steps: int
+    valid_steps: int
+    reporttypes: List[ReportTypes]
+    optimizer_kwargs: Dict[str, Any] = {"lr": 1e-3, "weight_decay": 1e-5}
+    scheduler_kwargs: Optional[Dict[str, Any]] = {"factor": 0.1, "patience": 10}
+    earlystop_kwargs: Optional[Dict[str, Any]] = {
+        "save": False,
+        "verbose": True,
+        "patience": 10,
+    }
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @field_validator("logdir")
+    @classmethod
+    def check_path(cls, logdir: Path) -> Path:  # noqa: N805
+        if isinstance(logdir, str):
+            logdir = Path(logdir)
+        if not logdir.resolve().exists():
+            logdir.mkdir(parents=True)
+            logger.info(f"Created logdir {logdir.absolute()}")
+        return logdir

{mltrainer-0.3.0 → mltrainer-0.3.1}/src/mltrainer/trainer.py

@@ -1,6 +1,8 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterator
 from datetime import datetime
 from pathlib import Path
-from typing import Callable, Dict, Iterator, Optional, Tuple, Type, TypeVar
+from typing import TYPE_CHECKING, TypeVar
 
 import mlflow
 import torch
@@ -11,10 +13,17 @@ from tqdm import tqdm
 
 from mltrainer import ReportTypes, TrainerSettings
 
+if TYPE_CHECKING:
+    from torch.utils.tensorboard.writer import SummaryWriter
+
 OptimizerType = TypeVar("OptimizerType", bound=torch.optim.Optimizer)
 
+# A batch input is either a single tensor or a tuple of tensors (e.g. packed
+# sequences); the target is always a tensor.
+BatchTensor = torch.Tensor | tuple[torch.Tensor, ...]
+
 
-def dir_add_timestamp(log_dir: Optional[Path] = None) -> Path:
+def dir_add_timestamp(log_dir: Path | None = None) -> Path:
     if log_dir is None:
         log_dir = Path(".")
     log_dir = Path(log_dir)
@@ -26,27 +35,110 @@ def dir_add_timestamp(log_dir: Optional[Path] = None) -> Path:
     return log_dir
 
 
+class Reporter(ABC):
+    """Reports training metrics to a single backend.
+
+    A new backend is added by subclassing (open for extension); the training
+    loop that iterates over reporters never changes (closed for modification).
+    """
+
+    @abstractmethod
+    def report(
+        self,
+        epoch: int,
+        train_loss: float,
+        test_loss: float,
+        metric_dict: dict[str, float],
+    ) -> None: ...
+
+
+class RayReporter(Reporter):
+    def __init__(self) -> None:
+        try:
+            import ray
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "ray is required for RAY reporting. "
+                "Install it with: uv add 'mltrainer[tune]'"
+            )
+        self._ray = ray
+
+    def report(
+        self,
+        epoch: int,
+        train_loss: float,
+        test_loss: float,
+        metric_dict: dict[str, float],
+    ) -> None:
+        self._ray.train.report(  # ty: ignore[possibly-missing-attribute]
+            {
+                "iterations": epoch,
+                "train_loss": train_loss,
+                "test_loss": test_loss,
+                **metric_dict,
+            }
+        )
+
+
+class MLflowReporter(Reporter):
+    def __init__(self, optimizer: Optimizer) -> None:
+        self.optimizer = optimizer
+
+    def report(
+        self,
+        epoch: int,
+        train_loss: float,
+        test_loss: float,
+        metric_dict: dict[str, float],
+    ) -> None:
+        mlflow.log_metric("Loss/train", train_loss, step=epoch)
+        mlflow.log_metric("Loss/test", test_loss, step=epoch)
+        for m in metric_dict:
+            mlflow.log_metric(f"metric/{m}", metric_dict[m], step=epoch)
+        lr = self.optimizer.param_groups[0]["lr"]
+        mlflow.log_metric("learning_rate", lr, step=epoch)
+
+
+class TensorBoardReporter(Reporter):
+    def __init__(self, writer: "SummaryWriter", optimizer: Optimizer) -> None:
+        self.writer = writer
+        self.optimizer = optimizer
+
+    def report(
+        self,
+        epoch: int,
+        train_loss: float,
+        test_loss: float,
+        metric_dict: dict[str, float],
+    ) -> None:
+        self.writer.add_scalar("Loss/train", train_loss, epoch)
+        self.writer.add_scalar("Loss/test", test_loss, epoch)
+        for m in metric_dict:
+            self.writer.add_scalar(f"metric/{m}", metric_dict[m], epoch)
+        lr = self.optimizer.param_groups[0]["lr"]
+        self.writer.add_scalar("learning_rate", lr, epoch)
+
+
 class Trainer:
     def __init__(
         self,
         model: torch.nn.Module,
         settings: TrainerSettings,
         loss_fn: Callable,
-        optimizer: Type[Optimizer],
+        optimizer: type[Optimizer],
         traindataloader: Iterator,
         validdataloader: Iterator,
-        scheduler: Optional[Callable],
-        device: Optional[str] = None,
+        scheduler: Callable | None,
+        device: str | None = None,
     ) -> None:
         self.model = model
         self.settings = settings
         self.log_dir = dir_add_timestamp(settings.logdir)
         self.loss_fn = loss_fn
-        self.optimizer = optimizer
         self.traindataloader = traindataloader
         self.validdataloader = validdataloader
         self.device = device
-        self.writer = None
+        self.writer: "SummaryWriter | None" = None
 
         if self.device:
             self.model.to(self.device)
@@ -66,7 +158,7 @@ class Trainer:
                 "Found earlystop_kwargs in settings."
                 "Set to None if you dont want earlystopping."
             )
-            self.early_stopping: Optional[EarlyStopping] = EarlyStopping(
+            self.early_stopping: EarlyStopping | None = EarlyStopping(
                 self.log_dir, **settings.earlystop_kwargs
             )
         else:
@@ -81,6 +173,34 @@ class Trainer:
             TOMLSerializer.save(model, self.log_dir / "model.toml")
             TOMLSerializer.save(settings, self.log_dir / "settings.toml")
 
+        self.reporters = self._build_reporters()
+
+    def _build_reporters(self) -> list[Reporter]:
+        """Map the requested report types onto reporter objects once, so the
+        per-epoch report loop stays closed for modification."""
+        reporters: list[Reporter] = []
+        if ReportTypes.RAY in self.settings.reporttypes:
+            reporters.append(RayReporter())
+        if ReportTypes.MLFLOW in self.settings.reporttypes:
+            reporters.append(MLflowReporter(self.optimizer))
+        if ReportTypes.TENSORBOARD in self.settings.reporttypes:
+            assert self.writer is not None
+            reporters.append(TensorBoardReporter(self.writer, self.optimizer))
+        return reporters
+
+    def _to_device(
+        self, x: BatchTensor, y: torch.Tensor
+    ) -> tuple[BatchTensor, torch.Tensor]:
+        """Move a batch onto self.device, handling tuple inputs (e.g. packed
+        sequences) as well as single tensors."""
+        if not self.device:
+            return x, y
+        if isinstance(x, tuple):
+            x = tuple(xx.to(self.device) for xx in x)
+        else:
+            x = x.to(self.device)
+        return x, y.to(self.device)
+
     def __del__(self):
         """Cleanup method to ensure proper resource handling"""
         if hasattr(self, "writer") and self.writer is not None:
@@ -115,38 +235,28 @@ class Trainer:
         train_steps = self.settings.train_steps
         for _ in tqdm(range(train_steps), colour="#1e4706"):
             x, y = next(iter(self.traindataloader))
-            if self.device:
-                if isinstance(x, tuple):
-                    x = tuple(xx.to(self.device) for xx in x)
-                    y = y.to(self.device)
-                else:
-                    x, y = x.to(self.device), y.to(self.device)
+            x, y = self._to_device(x, y)
             self.optimizer.zero_grad()
             yhat = self.model(x)
             loss = self.loss_fn(yhat, y)
             loss.backward()
-            self.optimizer.step()  # type: ignore
+            self.optimizer.step()
             train_loss += loss.cpu().detach().numpy()
         train_loss /= train_steps
 
         return train_loss
 
-    def evalbatches(self) -> Tuple[Dict[str, float], float]:
+    def evalbatches(self) -> tuple[dict[str, float], float]:
         """Evaluate model on validation data with proper device handling"""
         self.model.eval()
         valid_steps = self.settings.valid_steps
         test_loss: float = 0.0
-        metric_dict: Dict[str, float] = {}
+        metric_dict: dict[str, float] = {}
 
         with torch.no_grad():  # Prevent gradient computation during evaluation
             for _ in range(valid_steps):
                 x, y = next(iter(self.validdataloader))
-                if self.device:
-                    if isinstance(x, tuple):
-                        x = tuple(xx.to(self.device) for xx in x)
-                        y = y.to(self.device)
-                    else:
-                        x, y = x.to(self.device), y.to(self.device)
+                x, y = self._to_device(x, y)
 
                 # Forward pass
                 yhat = self.model(x)
@@ -161,7 +271,7 @@ class Trainer:
         # Average the results
         test_loss /= valid_steps
         for key in metric_dict:
-            metric_dict[str(key)] = metric_dict[str(key)] / valid_steps
+            metric_dict[key] = metric_dict[key] / valid_steps
 
         # Handle scheduler
         if self.scheduler:
@@ -172,54 +282,32 @@ class Trainer:
 
         return metric_dict, test_loss
 
-    def report(
-        self, epoch: int, train_loss: float, test_loss: float, metric_dict: Dict
-    ) -> None:
-        if (self.last_epoch != 0) and (epoch == 0):
+    def _resolve_epoch(self, epoch: int) -> int:
+        """Offset the epoch counter when resuming a previous training run."""
+        if self.last_epoch != 0 and epoch == 0:
             self.last_epoch += 1
             logger.info(f"Resuming epochs from previous training at {self.last_epoch}")
         if self.last_epoch != 0:
             epoch += self.last_epoch
-        reporttypes = self.settings.reporttypes
-        self.test_loss = test_loss
-
-        if ReportTypes.RAY in reporttypes:
-            try:
-                import ray
-            except ModuleNotFoundError:
-                raise ModuleNotFoundError(
-                    "ray is required for RAY reporting. "
-                    "Install it with: uv add 'mltrainer[tune]'"
-                )
-            ray.train.report(
-                {
-                    "iterations": epoch,
-                    "train_loss": train_loss,
-                    "test_loss": test_loss,
-                    **metric_dict,
-                }
-            )
+        return epoch
 
-        if ReportTypes.MLFLOW in reporttypes:
-            mlflow.log_metric("Loss/train", train_loss, step=epoch)
-            mlflow.log_metric("Loss/test", test_loss, step=epoch)
-            for m in metric_dict:
-                mlflow.log_metric(f"metric/{m}", metric_dict[m], step=epoch)
-            lr = [group["lr"] for group in self.optimizer.param_groups][0]  # type: ignore
-            mlflow.log_metric("learning_rate", lr, step=epoch)
+    def report(
+        self,
+        epoch: int,
+        train_loss: float,
+        test_loss: float,
+        metric_dict: dict[str, float],
+    ) -> None:
+        epoch = self._resolve_epoch(epoch)
+        self.test_loss = test_loss
 
-        if ReportTypes.TENSORBOARD in reporttypes:
-            assert self.writer is not None
-            self.writer.add_scalar("Loss/train", train_loss, epoch)
-            self.writer.add_scalar("Loss/test", test_loss, epoch)
-            for m in metric_dict:
-                self.writer.add_scalar(f"metric/{m}", metric_dict[m], epoch)
-            lr = [group["lr"] for group in self.optimizer.param_groups][0]  # type: ignore
-            self.writer.add_scalar("learning_rate", lr, epoch)
+        for reporter in self.reporters:
+            reporter.report(epoch, train_loss, test_loss, metric_dict)
 
         metric_scores = [f"{v:.4f}" for v in metric_dict.values()]
         logger.info(
-            f"Epoch {epoch} train {train_loss:.4f} test {test_loss:.4f} metric {metric_scores}"  # noqa E501
+            f"Epoch {epoch} train {train_loss:.4f} test {test_loss:.4f} "
+            f"metric {metric_scores}"
         )
 
 

{mltrainer-0.3.0 → mltrainer-0.3.1}/src/mltrainer/transformer.py

@@ -21,6 +21,9 @@ class PositionalEncoding(nn.Module):
     It adds positional information to the input embeddings.
     """
 
+    pe: Tensor  # registered as a buffer in __init__; annotated so type checkers
+    # know it is a Tensor rather than a generic Module attribute.
+
     def __init__(self, d_model: int, dropout: float = 0.1, max_seq_len: int = 5000):
         super().__init__()
         self.dropout = nn.Dropout(p=dropout)

mltrainer-0.3.0/Makefile

@@ -1,7 +0,0 @@
-format:
-	black mltrainer
-	isort mltrainer
-
-lint:
-	flake8 mltrainer
-	mypy mltrainer

mltrainer-0.3.0/src/mltrainer/settings.py

@@ -1,177 +0,0 @@
-from __future__ import annotations
-
-from enum import Enum
-from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional
-
-from loguru import logger
-from pydantic import BaseModel, ConfigDict, field_validator
-
-
-class FileTypes(Enum):
-    JPG = ".jpg"
-    PNG = ".png"
-    TXT = ".txt"
-    ZIP = ".zip"
-    TGZ = ".tgz"
-    TAR = ".tar.gz"
-    GZ = ".gz"
-
-
-class ReportTypes(Enum):
-    TOML = "TOML"
-    TENSORBOARD = "TENSORBOARD"
-    MLFLOW = "MLFLOW"
-    RAY = "RAY"
-
-
-class FormattedBase(BaseModel):
-    def __str__(self) -> str:
-        return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
-
-    def __repr__(self) -> str:
-        return "\n".join(f"{k}: {v}" for k, v in self.__dict__.items())
-
-
-class TrainerSettings(FormattedBase):
-    epochs: int
-    metrics: List[Callable]
-    logdir: Path
-    train_steps: int
-    valid_steps: int
-    reporttypes: List[ReportTypes]
-    optimizer_kwargs: Dict[str, Any] = {"lr": 1e-3, "weight_decay": 1e-5}
-    scheduler_kwargs: Optional[Dict[str, Any]] = {"factor": 0.1, "patience": 10}
-    earlystop_kwargs: Optional[Dict[str, Any]] = {
-        "save": False,
-        "verbose": True,
-        "patience": 10,
-    }
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-
-    @field_validator("logdir")
-    @classmethod
-    def check_path(cls, logdir: Path) -> Path:  # noqa: N805
-        if isinstance(logdir, str):
-            logdir = Path(logdir)
-        if not logdir.resolve().exists():
-            logdir.mkdir(parents=True)
-            logger.info(f"Created logdir {logdir.absolute()}")
-        return logdir
-
-    # def to_toml_dict(self) -> dict:
-    #     """Convert the model to a TOML-compatible dictionary."""
-    #     data = self.model_dump(exclude_none=True)
-    #
-    #     # Handle special cases
-    #     if "metrics" in data:
-    #         data["metrics"] = [
-    #             f"{metric.__module__}.{metric.__name__}"
-    #             if isinstance(metric, Callable) and hasattr(metric, "__name__")
-    #             else f"{metric.__class__.__module__}.{metric.__class__.__name__}"
-    #             for metric in data["metrics"]
-    #         ]
-    #
-    #     # Convert Path objects to strings
-    #     for key, value in data.items():
-    #         if isinstance(value, Path):
-    #             data[key] = str(value)
-    #
-    #     return data
-    #
-    # def save_toml_example(self, file_path: str | Path) -> None:
-    #     """Save an example TOML configuration file with comments."""
-    #     if isinstance(file_path, str):
-    #         file_path = Path(file_path)
-    #
-    #     data = self.to_toml_dict()
-    #     lines = []
-    #
-    #     # Start with just the simple scalar values we know work
-    #     lines.append("epochs = {}\n".format(data["epochs"]))
-    #     metric_paths = [
-    #         f"{metric.__module__}.{metric.__name__}" for metric in self.metrics
-    #     ]
-    #     lines.append("metrics = {}\n".format(str(metric_paths).replace("'", '"')))
-    #     lines.append('logdir = "{}"\n'.format(str(data["logdir"])))
-    #     lines.append("train_steps = {}\n".format(data["train_steps"]))
-    #     lines.append("valid_steps = {}\n".format(data["valid_steps"]))
-    #
-    #     # Add reporttypes as list of strings (enum names)
-    #     report_types = [rt.value for rt in data["reporttypes"]]
-    #     lines.append("reporttypes = {}\n".format(str(report_types).replace("'", '"')))
-    #
-    #     if "optimizer_kwargs" in data:
-    #         lines.append(
-    #             self.dict_to_toml("optimizer_kwargs", data["optimizer_kwargs"])
-    #         )
-    #
-    #     if "scheduler_kwargs" in data:
-    #         lines.append(
-    #             self.dict_to_toml("scheduler_kwargs", data["scheduler_kwargs"])
-    #         )
-    #
-    #     if "earlystop_kwargs" in data:
-    #         lines.append(
-    #             self.dict_to_toml("earlystop_kwargs", data["earlystop_kwargs"])
-    #         )
-    #
-    #     with open(file_path, "w", encoding="utf-8") as f:
-    #         f.write("".join(lines))
-    #
-    # @staticmethod
-    # def dict_to_toml(name: str, data: dict) -> str:
-    #     """
-    #     Convert a dictionary into a TOML table format.
-    #
-    #     Args:
-    #         name (str): Name of the TOML section.
-    #         data (dict): Dictionary to convert.
-    #
-    #     Returns:
-    #         str: TOML formatted table string.
-    #     """
-    #     toml_lines = [f"[{name}]"]
-    #     for k, v in data.items():
-    #         # Convert booleans to lowercase ('true'/'false')
-    #         if isinstance(v, bool):
-    #             v = str(v).lower()
-    #         toml_lines.append(f"{k} = {v}")
-    #     return "\n".join(toml_lines) + "\n"
-    #
-    # @classmethod
-    # def from_toml(cls, file_path: str | Path) -> "TrainerSettings":
-    #     """Load the model from a TOML file."""
-    #     if isinstance(file_path, str):
-    #         file_path = Path(file_path)
-    #
-    #     with open(file_path, "rb") as f:
-    #         data = tomllib.load(f)
-    #
-    #     print(data)
-    #
-    #     # Convert string paths to Path objects
-    #     if "logdir" in data and isinstance(data["logdir"], str):
-    #         data["logdir"] = Path(data["logdir"])
-    #         # Convert reporttype strings back to enums
-    #     if "reporttypes" in data:
-    #         data["reporttypes"] = [ReportTypes(rt) for rt in data["reporttypes"]]
-    #
-    #     # Import metrics from their module paths
-    #     if "metrics" in data:
-    #         metrics = []
-    #         for metric_path in data["metrics"]:
-    #             try:
-    #                 module_path, func_name = metric_path.rsplit(".", 1)
-    #                 module = __import__(module_path, fromlist=[func_name])
-    #                 metric = getattr(module, func_name)
-    #                 metrics.append(metric)
-    #             except (ImportError, AttributeError) as e:
-    #                 logger.warning(
-    #                     f"Could not load metric {metric_path}. "
-    #                     f"Error: {e}. Using placeholder."
-    #                 )
-    #         data["metrics"] = metrics
-    #
-    #     return cls(**data)