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 forward() / metrics() methods. Something like

  class Estimator[**P](nn.Module):
      def forward(
          self,
          input: Tensor,
          *args: P.args,
          **kwargs: P.kwargs,
      ) -> Tensor:
          ...
  
  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 (since ParamSpec isn't exactly controllable in that way; it's meant make a generic signature expressible as a type variable, not to express particular shapes of the types in that signature).

  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.

Typing planning

There are currently two package-wide typing strategies geared toward future Python version bumps and type system changes:

• Using Unpack with a type variable with a TypedDict upper bound, rather than with a TypedDict directly. This is commented on extensively in trainlib/estimator.py, but the short version is we want our key objects to remain generic in their method signatures when they are expected to operate on session-specific structured batch data. For instance, you can imagine an Estimator's forward method accepting a typed batch like

  def forward(self, batch: BatchKwargs): ...
  
  # or
  
  def forward(self, **kwargs: Unpack[Batchwargs]): ...
  

  These are both currently supported patterns. Problems arise, however, when we want an inherited class to specifically operate on an inherited type of BatchKwargs

  class BatchKwargs(TypedDict):
      inputs: Tensor
  
  class ExtraBatchKwargs(BatchKwargs):
      inputs: Tensor
      labels: Tensor
  
  class A(Estimator):
      def forward(self, batch: BatchKwargs): ...
  
  # want to require Estimator subclass handles both inputs and labels
  class B(A):
      def forward(self, batch: ExtraBatchKwargs): ...
  

  Note how ExtraBatchKwargs doesn't use NotRequired for the new labels attribute; the intention here is to actually require that B should handle more attributes than A. That breaks LSP, so this doesn't work. You'd need to either break out the Estimator hierarchies, or decide that B could have inputs typed as BatchKwargs and possibly implement some extra handling to check if batches of type ExtraBatchKwargs are passed in (note that passing in a batch of type ExtraBatchKwargs to A.forward is perfectly fine) and supply labels. Or... you make the class generic in the batch type, and during inheritance you can establish increasingly tight upper bounds on the batch type to accomplish precisely the desired effect. That is, you can formally require subclasses handle new batch attributes without breaking LSP. This could look like:

  class A[Kw: BatchKwargs](Estimator):
      def forward(self, batch: Kw): ...
  
  class B[Kw: ExtraBatchKwargs](A[Kw]):
      def forward(self, batch: Kw): ...
  

  This formally allows B to inherit from A while depending on new, non-optional attributes in its batch type; you can achieve that with refined type upper bound in a generic subtype. So this pattern works now, but what doesn't work is further unpacking Kw in this pattern into kwargs. That is,

  class A[Kw: BatchKwargs](Estimator):
      def forward(self, **kwargs: Unpack[Kw]): ...
  
  class B[Kw: ExtraBatchKwargs](A[Kw]):
      def forward(self, **kwargs: Unpack[Kw]): ...
  

  As sound as this is (and it is sound, described in PEP 827), Unpack isn't supported for type variables; it has to be used directly with TypedDict references. I have the latter in place across the package anticipating eventual support for this pattern, but things are subject to change. I may simply toggle back to the currently supported pattern provided it's very similar, and if/when PEP 827 is integrated, the change will be as simply as adding back the Unpack statements and keyword expansions.

• At the time of writing, trainlib attempts to remain compatible with Python 3.12. This means the convenient lazy evaluation of type annotations can't be fully embraced across the package, but we use __future__ where possible to make the transition easy when eventually bumping to 3.14+ (since the imports can just be removed at that time). For instance, class definitions that need to reference the class type itself shouldn't use string annotations but instead direct type references:

  # not
  class RecordIndex:
      @overload
      def __getitem__(self, key: slice) -> "RecordIndex": ...
  
  # but
  from __future__ import annotations
  
  class RecordIndex:
      @overload
      def __getitem__(self, key: slice) -> RecordIndex: ...