trainlib 0.3.0

Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training.

Overview

Minimal framework for ML modeling, supporting advanced dataset operations and streamlined training workflows.

Install

The trainlib package can be installed from PyPI:

pip install trainlib

Development

• Initialize/synchronize the project with uv sync, creating a virtual environment with base package dependencies.
• Depending on needs, install the development dependencies with uv sync --extra dev.

Testing

• To run the unit tests, make sure to first have the test dependencies installed with uv sync --extra test, then run make test.
• For notebook testing, run make install-kernel to make the environment available as a Jupyter kernel (to be selected when running notebooks).

Documentation

• Install the documentation dependencies with uv sync --extra doc.
• Run make docs-build (optionally preceded by make docs-clean), and serve locally with make docs-serve.

Development remarks

• Across Trainer / Estimator / Dataset, I've considered a ParamSpec-based typing scheme to better orchestrate alignment in the Trainer.train() loop, e.g., so we can statically check whether a dataset appears to be fulfilling the argument requirements for the estimator's loss() / metrics() methods. Something like

  class Estimator[**P](nn.Module):
      def loss(
          self,
          input: Tensor,
          *args: P.args,
          **kwargs: P.kwargs,
      ) -> Generator:
          ...
  
  class Trainer[**P]:
      def __init__(
          self,
          estimator: Estimator[P],
          ...
      ): ...
  

  might be how we begin threading signatures. But ensuring dataset items can match P is challenging. You can consider a "packed" object where we obfuscate passing data through P-signatures:

  class PackedItem[**P]:
      def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
          self._args = args
          self._kwargs = kwargs
  
      def apply[R](self, func: Callable[P, R]) -> R:
          return func(*self._args, **self._kwargs)
  
  
  class BatchedDataset[U, R, I, **P](Dataset):
      @abstractmethod
      def _process_item_data(
          self,
          item_data: I,
          item_index: int,
      ) -> PackedItem[P]:
          ...
  
      def __iter__(self) -> Iterator[PackedItem[P]]:
          ...
  

  Meaningfully shaping those signatures is what remains, but you can't really do this, not with typical type expression flexibility. For instance, if I'm trying to appropriately type my base TupleDataset:

  class SequenceDataset[I, **P](HomogenousDataset[int, I, I, P]):
      ...
  
  class TupleDataset[I](SequenceDataset[tuple[I, ...], "?"]):
      ...
  

  Here there's no way for me to shape a ParamSpec to indicate arbitrarily many arguments of a fixed type (I in this case) to allow me to unpack my item tuples into an appropriate PackedItem.

  Until this (among other issues) becomes clearer, I'm setting up around a simpler TypedDict type variable. We won't have particularly strong static checks for item alignment inside Trainer, but this seems about as good as I can get around the current infrastructure.