Metadata-Version: 2.4
Name: trainlib
Version: 0.6.1
Summary: Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training.
License-Expression: MIT
Project-URL: Homepage, https://doc.olog.io/trainlib
Project-URL: Documentation, https://doc.olog.io/trainlib
Project-URL: Repository, https://git.olog.io/olog/trainlib
Project-URL: Issues, https://git.olog.io/olog/trainlib/issues
Keywords: machine-learning
Classifier: Programming Language :: Python
Classifier: Operating System :: OS Independent
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: End Users/Desktop
Requires-Python: >=3.12
Description-Content-Type: text/markdown
Requires-Dist: torch
Requires-Dist: pandas
Requires-Dist: colorama>=0.4.6
Requires-Dist: matplotlib>=3.10.8
Requires-Dist: numpy>=2.4.1
Requires-Dist: tensorboard>=2.20.0
Requires-Dist: tqdm>=4.67.1
Requires-Dist: rich>=13.7
Requires-Dist: setuptools<=81.0.0
Provides-Extra: dev
Requires-Dist: ipykernel; extra == "dev"
Provides-Extra: doc
Requires-Dist: furo; extra == "doc"
Requires-Dist: myst-parser; extra == "doc"
Requires-Dist: sphinx; extra == "doc"
Requires-Dist: sphinx-togglebutton; extra == "doc"
Requires-Dist: sphinx-autodoc-typehints; extra == "doc"
Provides-Extra: test
Requires-Dist: pytest; extra == "test"

# Overview
Minimal framework for ML modeling, supporting advanced dataset operations and
streamlined training workflows.
    
# Install
The `trainlib` package can be installed from PyPI:

```sh
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

  ```py
  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:

  ```py
  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`:

  ```py
  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

  ```py
  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`

  ```py
  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:

  ```py
  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,

  ```py
  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:

  ```py
  # not
  class RecordIndex:
      @overload
      def __getitem__(self, key: slice) -> "RecordIndex": ...

  # but
  from __future__ import annotations

  class RecordIndex:
      @overload
      def __getitem__(self, key: slice) -> RecordIndex: ...
  ```
