fitstream 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

fitstream/__init__.py

@@ -0,0 +1,13 @@
+from .batching import iter_batches as iter_batches
+from .events import Event as Event
+from .fit import (
+    augment as augment,
+    early_stop as early_stop,
+    epoch_stream as epoch_stream,
+    pipe as pipe,
+    take as take,
+    tap as tap,
+    tick as tick,
+)
+from .sinks import collect as collect, collect_jsonl as collect_jsonl, collect_pd as collect_pd
+from .augmentations import validation_loss as validation_loss

fitstream/augmentations.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import torch
+from torch import nn
+
+
+def validation_loss(
+    val_data: Sequence[torch.Tensor],
+    loss_fn: Callable[..., torch.Tensor],
+    *,
+    key: str = "val_loss",
+    last_label: bool = True,
+) -> Callable[[dict[str, Any]], dict[str, float]]:
+    """Create an augmentation that computes validation loss.
+
+    Args:
+        val_data: Tuple of tensors with batch dimension first. When ``last_label=True``,
+            the last tensor is treated as the label tensor and all preceding tensors are
+            passed to the model.
+        loss_fn: Loss function called as ``loss_fn(pred, labels)`` when ``last_label=True``,
+            otherwise as ``loss_fn(pred)``.
+        key: Name of the key to store the computed loss under.
+        last_label: Whether the last tensor in ``val_data`` is the label tensor.
+
+    Notes:
+        Assumes the model and validation tensors are already on the same device.
+        Validation is computed on the full validation set (no batching).
+    """
+    if not val_data:
+        raise ValueError("val_data must contain at least one tensor.")
+    if last_label and len(val_data) < 2:
+        raise ValueError("last_label=True requires at least two tensors (inputs and labels).")
+
+    if last_label:
+        *inputs, labels = val_data
+    else:
+        inputs = list(val_data)
+        labels = None
+
+    def compute(event: dict[str, Any]) -> dict[str, float]:
+        model = event["model"]
+        if not isinstance(model, nn.Module):
+            raise TypeError("validation_loss expects an event containing a 'model' key.")
+
+        was_training = model.training
+        model.eval()
+        with torch.no_grad():
+            preds = model(*inputs)
+            loss = loss_fn(preds, labels) if last_label else loss_fn(preds)
+        if was_training:
+            model.train()
+        return {key: loss.detach().cpu().item()}
+
+    return compute

fitstream/batching.py

@@ -0,0 +1,37 @@
+from typing import Iterable, Sequence
+
+import torch
+
+
+def iter_batches(
+    *tensors, batch_size: int = 1, shuffle: bool = True, generator: torch.Generator | None = None
+) -> Iterable[Sequence[torch.Tensor]]:
+    """Yields batches from tensors, optionally shuffled.
+
+    Args:
+        *tensors: One or more tensors that share the same first dimension (sample axis). Each
+            yielded batch contains slices from each tensor aligned on that axis.
+        batch_size: Number of samples per batch. The final batch may be smaller if the sample
+            count is not divisible by the batch size.
+        shuffle: Whether to shuffle samples before batching. Shuffling uses the device of the
+            first tensor.
+        generator: Optional torch.Generator for deterministic shuffling.
+
+    Yields:
+        Tuples of tensors, one per input tensor, representing a batch.
+
+    Notes:
+        This function assumes all tensors have the same number of samples along dimension 0
+        and live on the same device. It does not perform explicit validation.
+    """
+    if not tensors:
+        return
+    if not shuffle:
+        tensor_batches = [tensor.split(batch_size) for tensor in tensors]
+        yield from zip(*tensor_batches)
+    else:
+        device = tensors[0].device
+        n_samples = tensors[0].shape[0]
+        idx = torch.randperm(n_samples, device=device, generator=generator)
+        for idx_chunk in idx.split(batch_size):
+            yield tuple(x[idx_chunk] for x in tensors)

fitstream/events.py

@@ -0,0 +1,19 @@
+from typing import TypedDict
+
+from torch import nn
+
+
+class Event(TypedDict):
+    """Per-epoch event emitted by fit/stream utilities.
+
+    Keys:
+        model: Live model reference updated each epoch.
+        step: 1-based epoch index.
+        train_loss: Mean training loss for the epoch.
+        train_time_sec: Wall-clock seconds spent in the epoch.
+    """
+
+    model: nn.Module
+    step: int
+    train_loss: float
+    train_time_sec: float

fitstream/fit.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator, Sequence
+from typing import Any
+import time
+
+import torch
+from torch import nn
+
+from .batching import iter_batches
+from .events import Event
+
+Transform = Callable[[Iterable[dict[str, Any]]], Iterable[dict[str, Any]]]
+
+
+def augment(
+    fn: Callable[[dict[str, Any]], dict[str, Any] | None],
+) -> Transform:
+    """Create a transform that merges extra keys into each event.
+
+    Args:
+        fn: Function called for each event. The returned mapping is shallow-merged
+            into the event. Returning ``None`` adds nothing.
+
+    Returns:
+        A transform that accepts an event stream and yields augmented events.
+    """
+
+    def transform(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        for event in events:
+            extra = fn(event) or {}
+            if not isinstance(extra, dict):
+                raise TypeError("augment function must return a dict or None.")
+            yield event | extra
+
+    return transform
+
+
+def pipe(stream: Iterable[dict[str, Any]], *stages: Transform) -> Iterable[dict[str, Any]]:
+    """Compose stream transforms left-to-right.
+
+    Args:
+        stream: Input event stream.
+        stages: Transform functions applied in order.
+
+    Returns:
+        The transformed event stream.
+    """
+    for stage in stages:
+        if not callable(stage):
+            raise TypeError("pipe stages must be callable.")
+        stream = stage(stream)
+    return stream
+
+
+def take(n: int) -> Transform:
+    """Limit an event stream to the first ``n`` events.
+
+    Can be used directly on a stream or as a pipe stage:
+
+    - ``take(10)(events)``
+    - ``pipe(events, take(10))``
+    """
+    if n < 0:
+        raise ValueError("n must be >= 0.")
+
+    def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        count = 0
+        for event in events:
+            if count >= n:
+                break
+            yield event
+            count += 1
+
+    return stage
+
+
+def tap(
+    fn: Callable[[dict[str, Any]], Any],
+) -> Transform:
+    """Create a stage that performs side effects and yields events unchanged."""
+    if not callable(fn):
+        raise TypeError("tap requires a callable.")
+
+    def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        for event in events:
+            fn(event)
+            yield event
+
+    return stage
+
+
+def tick(
+    fn: Callable[[], Any],
+) -> Transform:
+    """Create a stage that runs a no-arg callback per event and yields events unchanged."""
+    if not callable(fn):
+        raise TypeError("tick requires a callable.")
+
+    def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        for event in events:
+            fn()
+            yield event
+
+    return stage
+
+
+def early_stop(
+    key: str,
+    patience: int,
+) -> Transform:
+    """Yield events until the metric stops improving for `patience` steps.
+
+    Use as a pipe stage:
+
+    - ``pipe(events, early_stop(key="val_loss", patience=10))``
+    """
+    if patience < 1:
+        raise ValueError("patience must be >= 1.")
+
+    def apply(stream: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        best = float("inf")
+        bad = 0
+        for event in stream:
+            value = float(event[key])
+            if value < best:
+                best = value
+                bad = 0
+            else:
+                bad += 1
+            yield event
+            if bad >= patience:
+                break
+
+    return apply
+
+
+def epoch_stream(
+    train_data: Sequence[torch.Tensor],
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: Callable[..., torch.Tensor],
+    *,
+    batch_size: int = 1,
+    shuffle: bool = True,
+    last_label: bool = True,
+    generator: torch.Generator | None = None,
+) -> Iterator[Event]:
+    """Yield per-epoch training events from in-memory tensors.
+
+    Args:
+        train_data: Tuple of tensors with batch dimension first. When ``last_label=True``,
+            the last tensor is treated as the label tensor and all preceding tensors are
+            passed to the model.
+        model: PyTorch model to train.
+        optimizer: Optimizer instance constructed with the model parameters.
+        loss_fn: Loss function. Called as ``loss_fn(pred, labels)`` when ``last_label=True``,
+            otherwise as ``loss_fn(pred)``.
+        batch_size: Number of samples per batch.
+        shuffle: Whether to shuffle samples before batching.
+        last_label: Whether the last tensor in ``train_data`` is the label tensor.
+        generator: Optional torch.Generator forwarded to ``iter_batches`` for reproducible
+            shuffling.
+
+    Notes:
+        This function assumes the model and all tensors are already on the same device.
+        It does not copy tensors or take snapshots of model weights.
+    """
+    if batch_size <= 0:
+        raise ValueError("batch_size must be a positive integer.")
+    if not train_data:
+        raise ValueError("train_data must contain at least one tensor.")
+    if last_label and len(train_data) < 2:
+        raise ValueError("last_label=True requires at least two tensors (inputs and labels).")
+
+    step = 0
+    while True:
+        model.train()
+        epoch_start = time.perf_counter()
+        total_loss = torch.zeros((), device=train_data[0].device)
+        total_samples = 0
+
+        for batch in iter_batches(*train_data, batch_size=batch_size, shuffle=shuffle, generator=generator):
+            if last_label:
+                *inputs, labels = batch
+            else:
+                inputs = list(batch)
+                labels = None
+
+            preds = model(*inputs)
+            loss = loss_fn(preds, labels) if last_label else loss_fn(preds)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            batch_samples = inputs[0].shape[0]
+            total_loss += loss.detach() * batch_samples
+            total_samples += int(batch_samples)
+
+        step += 1
+        epoch_loss = (total_loss / total_samples).item()
+        yield Event(
+            model=model,
+            step=step,
+            train_loss=epoch_loss,
+            train_time_sec=time.perf_counter() - epoch_start,
+        )

fitstream/sinks.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Any, TextIO
+import json
+
+
+def _filter_event(
+    event: dict[str, Any],
+    *,
+    include: Sequence[str] | None,
+    exclude: Sequence[str] | None,
+) -> dict[str, Any]:
+    if include is not None and exclude is not None:
+        raise ValueError("Provide only one of include or exclude.")
+    if include is None:
+        excluded = {"model"}
+        if exclude is not None:
+            excluded.update(exclude)
+        return {k: v for k, v in event.items() if k not in excluded}
+    included = set(include)
+    included.discard("model")
+    return {k: event[k] for k in included if k in event}
+
+
+def collect(
+    events: Iterable[dict[str, Any]],
+    *,
+    include: Sequence[str] | None = None,
+    exclude: Sequence[str] | None = None,
+) -> list[dict[str, Any]]:
+    """Collect an event stream into a list of dicts."""
+    return [_filter_event(event, include=include, exclude=exclude) for event in events]
+
+
+def collect_jsonl(
+    events: Iterable[dict[str, Any]],
+    dest: str | Path | TextIO,
+    *,
+    include: Sequence[str] | None = None,
+    exclude: Sequence[str] | None = None,
+) -> None:
+    """Write events to JSONL (one JSON object per line)."""
+    if isinstance(dest, (str, Path)):
+        with Path(dest).open("w") as handle:
+            collect_jsonl(events, handle, include=include, exclude=exclude)
+        return
+    for event in events:
+        record = _filter_event(event, include=include, exclude=exclude)
+        dest.write(json.dumps(record) + "\n")
+
+
+def collect_pd(
+    events: Iterable[dict[str, Any]],
+    *,
+    include: Sequence[str] | None = None,
+    exclude: Sequence[str] | None = None,
+):
+    """Collect events into a pandas DataFrame."""
+    try:
+        import pandas as pd
+    except Exception as exc:  # pragma: no cover - depends on optional dependency
+        raise ImportError("pandas is required for collect_pd.") from exc
+    rows = [_filter_event(event, include=include, exclude=exclude) for event in events]
+    return pd.DataFrame(rows)

fitstream-0.1.1.dist-info/METADATA

@@ -0,0 +1,388 @@
+Metadata-Version: 2.3
+Name: fitstream
+Version: 0.1.1
+Summary: A library to train PyTorch models as a stream of events
+Keywords: machine-learning,deep-learning,neural-networks,pytorch,torch,training,metrics,pipeline,dataloader
+Author: Alex Shtoff
+Author-email: Alex Shtoff <alex.shtf@gmail.com>
+License:                                  Apache License
+                                    Version 2.0, January 2004
+                                 http://www.apache.org/licenses/
+         
+            TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+         
+            1. Definitions.
+         
+               "License" shall mean the terms and conditions for use, reproduction,
+               and distribution as defined by Sections 1 through 9 of this document.
+         
+               "Licensor" shall mean the copyright owner or entity authorized by
+               the copyright owner that is granting the License.
+         
+               "Legal Entity" shall mean the union of the acting entity and all
+               other entities that control, are controlled by, or are under common
+               control with that entity. For the purposes of this definition,
+               "control" means (i) the power, direct or indirect, to cause the
+               direction or management of such entity, whether by contract or
+               otherwise, or (ii) ownership of fifty percent (50%) or more of the
+               outstanding shares, or (iii) beneficial ownership of such entity.
+         
+               "You" (or "Your") shall mean an individual or Legal Entity
+               exercising permissions granted by this License.
+         
+               "Source" form shall mean the preferred form for making modifications,
+               including but not limited to software source code, documentation
+               source, and configuration files.
+         
+               "Object" form shall mean any form resulting from mechanical
+               transformation or translation of a Source form, including but
+               not limited to compiled object code, generated documentation,
+               and conversions to other media types.
+         
+               "Work" shall mean the work of authorship, whether in Source or
+               Object form, made available under the License, as indicated by a
+               copyright notice that is included in or attached to the work
+               (an example is provided in the Appendix below).
+         
+               "Derivative Works" shall mean any work, whether in Source or Object
+               form, that is based on (or derived from) the Work and for which the
+               editorial revisions, annotations, elaborations, or other modifications
+               represent, as a whole, an original work of authorship. For the purposes
+               of this License, Derivative Works shall not include works that remain
+               separable from, or merely link (or bind by name) to the interfaces of,
+               the Work and Derivative Works thereof.
+         
+               "Contribution" shall mean any work of authorship, including
+               the original version of the Work and any modifications or additions
+               to that Work or Derivative Works thereof, that is intentionally
+               submitted to Licensor for inclusion in the Work by the copyright owner
+               or by an individual or Legal Entity authorized to submit on behalf of
+               the copyright owner. For the purposes of this definition, "submitted"
+               means any form of electronic, verbal, or written communication sent
+               to the Licensor or its representatives, including but not limited to
+               communication on electronic mailing lists, source code control systems,
+               and issue tracking systems that are managed by, or on behalf of, the
+               Licensor for the purpose of discussing and improving the Work, but
+               excluding communication that is conspicuously marked or otherwise
+               designated in writing by the copyright owner as "Not a Contribution."
+         
+               "Contributor" shall mean Licensor and any individual or Legal Entity
+               on behalf of whom a Contribution has been received by Licensor and
+               subsequently incorporated within the Work.
+         
+            2. Grant of Copyright License. Subject to the terms and conditions of
+               this License, each Contributor hereby grants to You a perpetual,
+               worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+               copyright license to reproduce, prepare Derivative Works of,
+               publicly display, publicly perform, sublicense, and distribute the
+               Work and such Derivative Works in Source or Object form.
+         
+            3. Grant of Patent License. Subject to the terms and conditions of
+               this License, each Contributor hereby grants to You a perpetual,
+               worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+               (except as stated in this section) patent license to make, have made,
+               use, offer to sell, sell, import, and otherwise transfer the Work,
+               where such license applies only to those patent claims licensable
+               by such Contributor that are necessarily infringed by their
+               Contribution(s) alone or by combination of their Contribution(s)
+               with the Work to which such Contribution(s) was submitted. If You
+               institute patent litigation against any entity (including a
+               cross-claim or counterclaim in a lawsuit) alleging that the Work
+               or a Contribution incorporated within the Work constitutes direct
+               or contributory patent infringement, then any patent licenses
+               granted to You under this License for that Work shall terminate
+               as of the date such litigation is filed.
+         
+            4. Redistribution. You may reproduce and distribute copies of the
+               Work or Derivative Works thereof in any medium, with or without
+               modifications, and in Source or Object form, provided that You
+               meet the following conditions:
+         
+               (a) You must give any other recipients of the Work or
+                   Derivative Works a copy of this License; and
+         
+               (b) You must cause any modified files to carry prominent notices
+                   stating that You changed the files; and
+         
+               (c) You must retain, in the Source form of any Derivative Works
+                   that You distribute, all copyright, patent, trademark, and
+                   attribution notices from the Source form of the Work,
+                   excluding those notices that do not pertain to any part of
+                   the Derivative Works; and
+         
+               (d) If the Work includes a "NOTICE" text file as part of its
+                   distribution, then any Derivative Works that You distribute must
+                   include a readable copy of the attribution notices contained
+                   within such NOTICE file, excluding those notices that do not
+                   pertain to any part of the Derivative Works, in at least one
+                   of the following places: within a NOTICE text file distributed
+                   as part of the Derivative Works; within the Source form or
+                   documentation, if provided along with the Derivative Works; or,
+                   within a display generated by the Derivative Works, if and
+                   wherever such third-party notices normally appear. The contents
+                   of the NOTICE file are for informational purposes only and
+                   do not modify the License. You may add Your own attribution
+                   notices within Derivative Works that You distribute, alongside
+                   or as an addendum to the NOTICE text from the Work, provided
+                   that such additional attribution notices cannot be construed
+                   as modifying the License.
+         
+               You may add Your own copyright statement to Your modifications and
+               may provide additional or different license terms and conditions
+               for use, reproduction, or distribution of Your modifications, or
+               for any such Derivative Works as a whole, provided Your use,
+               reproduction, and distribution of the Work otherwise complies with
+               the conditions stated in this License.
+         
+            5. Submission of Contributions. Unless You explicitly state otherwise,
+               any Contribution intentionally submitted for inclusion in the Work
+               by You to the Licensor shall be under the terms and conditions of
+               this License, without any additional terms or conditions.
+               Notwithstanding the above, nothing herein shall supersede or modify
+               the terms of any separate license agreement you may have executed
+               with Licensor regarding such Contributions.
+         
+            6. Trademarks. This License does not grant permission to use the trade
+               names, trademarks, service marks, or product names of the Licensor,
+               except as required for reasonable and customary use in describing the
+               origin of the Work and reproducing the content of the NOTICE file.
+         
+            7. Disclaimer of Warranty. Unless required by applicable law or
+               agreed to in writing, Licensor provides the Work (and each
+               Contributor provides its Contributions) on an "AS IS" BASIS,
+               WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+               implied, including, without limitation, any warranties or conditions
+               of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+               PARTICULAR PURPOSE. You are solely responsible for determining the
+               appropriateness of using or redistributing the Work and assume any
+               risks associated with Your exercise of permissions under this License.
+         
+            8. Limitation of Liability. In no event and under no legal theory,
+               whether in tort (including negligence), contract, or otherwise,
+               unless required by applicable law (such as deliberate and grossly
+               negligent acts) or agreed to in writing, shall any Contributor be
+               liable to You for damages, including any direct, indirect, special,
+               incidental, or consequential damages of any character arising as a
+               result of this License or out of the use or inability to use the
+               Work (including but not limited to damages for loss of goodwill,
+               work stoppage, computer failure or malfunction, or any and all
+               other commercial damages or losses), even if such Contributor
+               has been advised of the possibility of such damages.
+         
+            9. Accepting Warranty or Additional Liability. While redistributing
+               the Work or Derivative Works thereof, You may choose to offer,
+               and charge a fee for, acceptance of support, warranty, indemnity,
+               or other liability obligations and/or rights consistent with this
+               License. However, in accepting such obligations, You may act only
+               on Your own behalf and on Your sole responsibility, not on behalf
+               of any other Contributor, and only if You agree to indemnify,
+               defend, and hold each Contributor harmless for any liability
+               incurred by, or claims asserted against, such Contributor by reason
+               of your accepting any such warranty or additional liability.
+         
+            END OF TERMS AND CONDITIONS
+         
+            APPENDIX: How to apply the Apache License to your work.
+         
+               To apply the Apache License to your work, attach the following
+               boilerplate notice, with the fields enclosed by brackets "[]"
+               replaced with your own identifying information. (Don't include
+               the brackets!)  The text should be enclosed in the appropriate
+               comment syntax for the file format. We also recommend that a
+               file or class name and description of purpose be included on the
+               same "printed page" as the copyright notice for easier
+               identification within third-party archives.
+         
+            Copyright [yyyy] [name of copyright owner]
+         
+            Licensed under the Apache License, Version 2.0 (the "License");
+            you may not use this file except in compliance with the License.
+            You may obtain a copy of the License at
+         
+                http://www.apache.org/licenses/LICENSE-2.0
+         
+            Unless required by applicable law or agreed to in writing, software
+            distributed under the License is distributed on an "AS IS" BASIS,
+            WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+            See the License for the specific language governing permissions and
+            limitations under the License.
+Classifier: Development Status :: 3 - Alpha
+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.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries
+Requires-Dist: torch>=2.6.0
+Requires-Python: >=3.12
+Project-URL: Documentation, https://fitstream.readthedocs.io/
+Project-URL: Homepage, https://github.com/alexshtf/fitstream
+Project-URL: Repository, https://github.com/alexshtf/fitstream
+Project-URL: Issues, https://github.com/alexshtf/fitstream/issues
+Description-Content-Type: text/markdown
+
+# FitStream
+A tiny library to make PyTorch experiment easy for small models and in-memory datasets.
+
+# Getting started
+Using `uv`
+```shell
+uv add fitstream
+```
+
+Using `pip`:
+```shell
+pip install fitstream
+```
+
+Training a model:
+```python
+from torch.optim import Adam
+
+from fitstream import epoch_stream, take # epoch_stream is the main entry point
+
+X, y = get_data()
+model = get_model()
+loss = get_loss()
+optimizer = Adam(model.parameters())
+
+# an infinite stream of training epochs (limit it with `take` or `early_stop`)
+events = epoch_stream((X, y), model, optimizer, loss, batch_size=32, shuffle=True)
+for event in take(10)(events):
+    print(f"step={event['step']}, loss={event['train_loss']}")
+# epoch=1, loss=...
+# epoch=2, loss=...
+# ...
+```
+
+# Basics
+The core idea of the library is "training loop as a stream of events".  The `epoch_stream` is just an iterable over 
+dictionaries comprising of the epoch, the model, and the training loss. Everything we do is transforming or enriching
+these events. FitStream provides a small `pipe(...)` helper to compose transformations left-to-right.
+
+## Augmentation
+The `augment` function turns an "augmenter" (a function that looks at an event and returns extra keys) into a stream
+transform stage. We typically compose stages with `pipe(...)`.
+
+Here is an example - we add the norm of the model parameters to each event:
+```python
+from torch import nn, linalg
+from fitstream import epoch_stream, augment, pipe
+
+def model_param_norm(ev: dict) -> dict:
+    model_params = nn.utils.parameters_to_vector(ev['model'].parameters())
+    return {'model_param_norm': linalg.norm(model_params)}
+
+
+events = pipe(
+    epoch_stream(...),
+    augment(model_param_norm),
+)
+for event in events:
+    print(f"step={event['step']}", 
+          f"model_param_norm={event['model_param_norm']}"
+    )
+```
+
+We also have some built-in augmentation functions. Here is an example of adding validation loss to each event:
+```python
+from torch import nn
+from fitstream import epoch_stream, augment, pipe, validation_loss
+
+validation_set = get_validation_set()
+events = pipe(
+    epoch_stream(...),
+    augment(validation_loss(validation_set, nn.CrossEntropyLoss())),
+)
+for event in events:
+    print(f"step={event['step']}, val_loss={event['val_loss']}")
+```
+
+We can, of course, augment the stream more than once:
+```python
+events = pipe(
+    epoch_stream(...),
+    augment(validation_loss(...)),
+    augment(model_param_norm),
+)
+for event in events:
+    print(f"step={event['step']}", 
+          f"val_loss={event['val_loss']}",
+          f"model_param_norm={event['model_param_norm']}"
+    )
+```
+
+## Selecting events
+Since the training loop is a standard Python iterable, you can use any Python selection logic. FitStream includes a
+small helper, `take(...)`, to limit the number of epochs:
+```python
+from fitstream import epoch_stream, take
+
+for event in take(100)(epoch_stream(...)):
+    print(event)
+# {'step': 1, ....}
+# {'step': 2, ...}
+# ...
+# { 'step': 100, ...}
+```
+
+`fitstream` has some of its own selection primitives, such as early stopping:
+```python
+from fitstream import augment, early_stop, epoch_stream, pipe, take, validation_loss
+
+events = pipe(
+    epoch_stream(...),
+    augment(validation_loss(...)),
+    take(500),  # safety cap
+    early_stop(key="val_loss", patience=10),
+)
+for event in events:
+    print(event)
+```
+
+## Side effects
+Sometimes you want to log metrics (or write to an external system) without changing the stream. Use `tap(fn)`:
+```python
+from fitstream import epoch_stream, pipe, tap, take
+
+events = pipe(
+    epoch_stream(...),
+    tap(lambda ev: print(ev["step"], ev["train_loss"])),
+    take(10),
+)
+list(events)
+```
+
+## Sinks
+Iterating over events and doing something yourself can be tedious, so we have some utilities to help you process the
+event stream.
+
+It is typically useful to collect all events into a list, but exclude the `model` and keep just the metrics. We have 
+the `collect` sink for that:
+```python
+from fitstream import collect, epoch_stream, take
+
+# collect 100 epochs to a list
+history = collect(take(100)(epoch_stream(...)))
+```
+
+We can also store them to a `jsonl` file:
+```python
+from fitstream import collect_jsonl, epoch_stream, take
+
+# collect 100 epochs to json
+collect_jsonl(take(100)(epoch_stream(...)), 'runs/my_experiment.jsonl')
+```
+
+# Documentation
+Full documentation is available at [https://fitstream.readthedocs.io/](https://fitstream.readthedocs.io/).
+
+# Development
+- After cloning this repo, run `make setup` to create a virtual environment and install all dependencies.
+- Building is done via `uv build`.
+- Running tests is done via `make test`
+- Building documentation via `make doc`
+- Linting via `make lint`

fitstream-0.1.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+fitstream/__init__.py,sha256=i4FTMn5gOoRnZhXxcGuRCIfGbIk_CQ3NMjvleebF6U8,423
+fitstream/augmentations.py,sha256=KeslNnDEOmmeXkc4d70hguWhl_jdFcEzvrPnmzuM00E,1991
+fitstream/batching.py,sha256=ECzf9A3YaZ2_AySZuSYl57fvqHH_DMZKwUIi6R7mGLs,1500
+fitstream/events.py,sha256=YpmHViJR8WHLK4vpBEPiG6avNM8y9tTm7Sx0DzVsiMM,444
+fitstream/fit.py,sha256=utXohEn5Yy-KSt_LEHx-vtsBIt8ehj_UHnPKRRKmoJE,6264
+fitstream/sinks.py,sha256=ZTsyyGam56pDMRVU6nSh-6lpWEX2amDYr3bYQwyxs9o,2117
+fitstream-0.1.1.dist-info/WHEEL,sha256=5DEXXimM34_d4Gx1AuF9ysMr1_maoEtGKjaILM3s4w4,80
+fitstream-0.1.1.dist-info/METADATA,sha256=ws_-6hbakxiYnOhGl2Mcf8vgVCoy6T89CBm7WY7uKqY,19049
+fitstream-0.1.1.dist-info/RECORD,,

{fitstream-0.1.0.dist-info → fitstream-0.1.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: uv 0.9.28
+Generator: uv 0.9.29
 Root-Is-Purelib: true
 Tag: py3-none-any

fitstream-0.1.0.dist-info/METADATA

@@ -1,5 +0,0 @@
-Metadata-Version: 2.3
-Name: fitstream
-Version: 0.1.0
-Summary: Add your description here
-Requires-Python: >=3.12

fitstream-0.1.0.dist-info/RECORD

@@ -1,4 +0,0 @@
-fitstream/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-fitstream-0.1.0.dist-info/WHEEL,sha256=fAguSjoiATBe7TNBkJwOjyL1Tt4wwiaQGtNtjRPNMQA,80
-fitstream-0.1.0.dist-info/METADATA,sha256=-IzjLUWYXirg121U3bxUSf1FWiqMJPfwLLYGW4cKTEM,112
-fitstream-0.1.0.dist-info/RECORD,,