Merge pull request #3 from alexshtf/ema_tap

Add ema stage,  tap frequency setting, and logging helpers

Author: alexshtf

README.md

@@ -130,14 +130,15 @@ for event in events:
 `min_delta` to ignore tiny noisy changes.
 
 ## Side effects
-Sometimes you want to log metrics (or write to an external system) without changing the stream. Use `tap(fn)`:
+Sometimes you want to log metrics (or write to an external system) without changing the stream.
+Use `tap(fn, every=...)` and the built-in `print_keys(...)` helper:
 ```python
-from fitstream import epoch_stream, pipe, tap, take
+from fitstream import epoch_stream, pipe, print_keys, tap, take
 
 events = pipe(
     epoch_stream(...),
-    tap(lambda ev: print(ev["step"], ev["train_loss"])),
-    take(10),
+    tap(print_keys("train_loss"), every=5),
+    take(20),
 )
 list(events)
 ```

docs/tutorial.md

@@ -344,20 +344,20 @@ across events:
 
 ### 7.1 Quick side effects with `tap(...)`
 
-FitStream includes a small helper for side effects: `tap(fn)` calls `fn(event)` and yields the event unchanged.
-It’s perfect for lightweight logging or writing metrics to an external system.
+FitStream includes a small helper for side effects: `tap(fn, every=...)` calls `fn(event)` every N events and yields
+the event unchanged. It’s perfect for lightweight logging or writing metrics to an external system.
 
 ```python
-from fitstream import augment, epoch_stream, pipe, take, tap, validation_loss
+from fitstream import augment, epoch_stream, pipe, print_keys, take, tap, validation_loss
 
 events = pipe(
     epoch_stream((x_train, y_train), model, optimizer, loss_fn, batch_size=512, shuffle=True),
     augment(validation_loss((x_val, y_val), loss_fn)),
-    tap(lambda event: print(f"epoch={event['step']:03d} val_loss={event['val_loss']:.4f}")),
+    tap(print_keys("train_loss", "val_loss"), every=5),
 )
 
 # Consume a few events to actually run it (streams are lazy).
-list(take(5)(events))
+list(take(15)(events))
 ```
 
 ### 7.2 Learning rate scheduling with `tick(...)`
@@ -390,48 +390,49 @@ If your scheduler needs a metric (e.g. `ReduceLROnPlateau`), use `tap(...)` inst
 
 ### 7.3 Example: exponential moving average (EMA)
 
-This stage adds a new key like `val_loss_ema` to each event.
+FitStream includes an `ema(...)` stage that adds a new key like `val_loss_ema` to each event.
 
 ```python
-from collections.abc import Iterable
-from typing import Any
-
-def ema(key: str, *, alpha: float = 0.2, out_key: str | None = None):
-    if not (0.0 < alpha <= 1.0):
-        raise ValueError("alpha must be in (0, 1].")
-    out_key = out_key or f"{key}_ema"
-
-    def stage(events: Iterable[dict[str, Any]]):
-        value_ema: float | None = None
-        for event in events:
-            value = float(event[key])
-            value_ema = value if value_ema is None else (1.0 - alpha) * value_ema + alpha * value
-            yield event | {out_key: value_ema}
-
-    return stage
+from fitstream import augment, ema, epoch_stream, pipe
+
+# Coefficient form: m = decay * m + (1 - decay) * x
+events = pipe(
+    epoch_stream(...),
+    augment(...),
+    ema("val_loss", decay=0.9),
+)
+
+# Half-life form (more intuitive tuning in "events until ~50% influence")
+events = pipe(
+    epoch_stream(...),
+    augment(...),
+    ema("val_loss", half_life=10),
+)
 ```
 
-### 7.4 Example: print progress every N epochs
+`ema(..., bias_correction=True)` is the default (Adam-style correction). You can disable it:
 
 ```python
-from collections.abc import Iterable
-from typing import Any
-
-def print_every(n: int, *, keys: tuple[str, ...] = ("train_loss", "val_loss")):
-    if n <= 0:
-        raise ValueError("n must be >= 1.")
-
-    def stage(events: Iterable[dict[str, Any]]):
-        for event in events:
-            if int(event["step"]) % n == 0:
-                parts = [f"epoch={event['step']:04d}"]
-                for k in keys:
-                    if k in event:
-                        parts.append(f"{k}={float(event[k]):.4f}")
-                print(" ".join(parts))
-            yield event
-
-    return stage
+from fitstream import augment, ema, epoch_stream, pipe
+
+events = pipe(
+    epoch_stream(...),
+    augment(...),
+    ema("val_loss", half_life=10, bias_correction=False),
+)
+```
+
+### 7.4 Combine smoothing + periodic logging
+
+```python
+from fitstream import augment, ema, epoch_stream, pipe, print_keys, tap
+
+events = pipe(
+    epoch_stream(...),
+    augment(...),
+    ema("val_loss", half_life=10),
+    tap(print_keys("train_loss", "val_loss", "val_loss_ema"), every=10),
+)
 ```
 
 ## 8) The “zero → hero” pipeline (put it all together)
@@ -453,7 +454,18 @@ from pathlib import Path
 import torch
 from torch import nn
 
-from fitstream import augment, collect_jsonl, early_stop, epoch_stream, pipe, take, validation_loss
+from fitstream import (
+    augment,
+    collect_jsonl,
+    early_stop,
+    ema,
+    epoch_stream,
+    pipe,
+    print_keys,
+    take,
+    tap,
+    validation_loss,
+)
 
 RUNS_DIR = Path("runs")
 RUNS_DIR.mkdir(exist_ok=True)
@@ -472,8 +484,8 @@ events = pipe(
     epoch_stream((x_train, y_train), model, optimizer, loss_fn, batch_size=512, shuffle=True),
     augment(validation_loss((x_val, y_val), loss_fn)),
     augment(model_param_norm),
-    ema("val_loss", alpha=0.2),
-    print_every(10, keys=("train_loss", "val_loss", "val_loss_ema", "param_l2")),
+    ema("val_loss", half_life=10),
+    tap(print_keys("train_loss", "val_loss", "val_loss_ema", "param_l2"), every=10),
     take(500),
     early_stop(key="val_loss", patience=20, mode="min", min_delta=1e-4),
 )

src/fitstream/__init__.py

@@ -2,9 +2,11 @@
 from .events import Event as Event
 from .fit import (
     augment as augment,
+    ema as ema,
     early_stop as early_stop,
     epoch_stream as epoch_stream,
     pipe as pipe,
+    print_keys as print_keys,
     take as take,
     tap as tap,
     tick as tick,

src/fitstream/fit.py

@@ -77,19 +77,80 @@ def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
 
 def tap(
     fn: Callable[[dict[str, Any]], Any],
+    *,
+    every: int = 1,
+    start: int = 1,
 ) -> Transform:
-    """Create a stage that performs side effects and yields events unchanged."""
+    """Create a stage that performs side effects and yields events unchanged.
+
+    Args:
+        fn: Callback applied to selected events.
+        every: Call ``fn`` every N events (event-count based).
+        start: 1-based event index at which callback scheduling starts.
+    """
     if not callable(fn):
         raise TypeError("tap requires a callable.")
+    if every < 1:
+        raise ValueError("every must be >= 1.")
+    if start < 1:
+        raise ValueError("start must be >= 1.")
 
     def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
-        for event in events:
-            fn(event)
+        for index, event in enumerate(events, start=1):
+            if index >= start and (index - start) % every == 0:
+                fn(event)
             yield event
 
     return stage
 
 
+def print_keys(
+    *keys: str,
+    precision: int = 4,
+    include_step: bool = True,
+    step_key: str = "step",
+) -> Callable[[dict[str, Any]], None]:
+    """Create an event callback that prints selected keys in one compact line.
+
+    Args:
+        *keys: Event keys to print.
+        precision: Number of digits after the decimal for numeric values.
+        include_step: Whether to include ``step_key`` first when present.
+        step_key: Event key used for the step prefix.
+    """
+    if precision < 0:
+        raise ValueError("precision must be >= 0.")
+    if not keys and not include_step:
+        raise ValueError("Provide at least one key when include_step=False.")
+
+    def format_value(value: Any) -> str:
+        match value:
+            case torch.Tensor() as tensor if tensor.numel() == 1:
+                return f"{float(tensor.detach().cpu().item()):.{precision}f}"
+            case bool() as boolean:
+                return str(boolean)
+            case int() | float() as number:
+                return f"{float(number):.{precision}f}"
+            case _:
+                return str(value)
+
+    def callback(event: dict[str, Any]) -> None:
+        parts: list[str] = []
+        if include_step and step_key in event:
+            try:
+                parts.append(f"{step_key}={int(event[step_key]):04d}")
+            except Exception:
+                parts.append(f"{step_key}={event[step_key]}")
+        for key in keys:
+            if key in event:
+                parts.append(f"{key}={format_value(event[key])}")
+            else:
+                parts.append(f"{key}=NA")
+        print(" ".join(parts))
+
+    return callback
+
+
 def tick(
     fn: Callable[[], Any],
 ) -> Transform:
@@ -105,6 +166,47 @@ def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
     return stage
 
 
+def ema(
+    key: str,
+    *,
+    decay: float | None = None,
+    half_life: float | None = None,
+    out_key: str | None = None,
+    bias_correction: bool = True,
+) -> Transform:
+    """Create a stage that adds an exponential moving average of ``key``.
+
+    Exactly one of ``decay`` or ``half_life`` must be provided.
+    The update rule is ``m = decay * m + (1 - decay) * x`` with ``m`` initialized to 0.
+    """
+    if (decay is None) == (half_life is None):
+        raise ValueError("Provide exactly one of decay or half_life.")
+    if half_life is not None:
+        if half_life <= 0.0:
+            raise ValueError("half_life must be > 0.")
+        decay = 2.0 ** (-1.0 / half_life)
+    assert decay is not None
+    if not (0.0 < decay < 1.0):
+        raise ValueError("decay must be in (0, 1).")
+
+    output_key = out_key or f"{key}_ema"
+
+    def stage(events: Iterable[dict[str, Any]]) -> Iterable[dict[str, Any]]:
+        aggregate = 0.0
+        t = 0
+        for event in events:
+            t += 1
+            value = float(event[key])
+            aggregate = decay * aggregate + (1.0 - decay) * value
+            if bias_correction:
+                smoothed = aggregate / (1.0 - (decay**t))
+            else:
+                smoothed = aggregate
+            yield event | {output_key: smoothed}
+
+    return stage
+
+
 def early_stop(
     key: str,
     patience: int,

tests/test_ema.py

@@ -0,0 +1,84 @@
+import pytest
+
+from fitstream import ema, pipe
+
+
+def test_ema_bias_correction_with_decay() -> None:
+    events = [{"loss": 10.0}, {"loss": 20.0}, {"loss": 30.0}]
+
+    result = list(ema("loss", decay=0.5)(events))
+
+    observed = [event["loss_ema"] for event in result]
+    expected = [10.0, 16.6666666667, 24.2857142857]
+    assert observed == pytest.approx(expected)
+
+
+def test_ema_can_disable_bias_correction() -> None:
+    events = [{"loss": 10.0}, {"loss": 20.0}, {"loss": 30.0}]
+
+    result = list(ema("loss", decay=0.5, bias_correction=False)(events))
+
+    observed = [event["loss_ema"] for event in result]
+    expected = [5.0, 12.5, 21.25]
+    assert observed == pytest.approx(expected)
+
+
+def test_ema_half_life_matches_equivalent_decay() -> None:
+    events = [{"loss": 10.0}, {"loss": 20.0}, {"loss": 30.0}]
+
+    from_decay = list(ema("loss", decay=0.5)(events))
+    from_half_life = list(ema("loss", half_life=1.0)(events))
+
+    observed_decay = [event["loss_ema"] for event in from_decay]
+    observed_half_life = [event["loss_ema"] for event in from_half_life]
+    assert observed_half_life == pytest.approx(observed_decay)
+
+
+def test_ema_supports_custom_output_key() -> None:
+    events = [{"loss": 10.0}]
+
+    result = list(ema("loss", decay=0.5, out_key="smooth_loss")(events))
+
+    assert "smooth_loss" in result[0]
+    assert "loss_ema" not in result[0]
+
+
+def test_ema_rejects_missing_decay_configuration() -> None:
+    with pytest.raises(ValueError):
+        ema("loss")
+
+
+def test_ema_rejects_conflicting_decay_configuration() -> None:
+    with pytest.raises(ValueError):
+        ema("loss", decay=0.9, half_life=10.0)
+
+
+@pytest.mark.parametrize("decay", [-0.1, 0.0, 1.0, 1.1])
+def test_ema_rejects_invalid_decay(decay: float) -> None:
+    with pytest.raises(ValueError):
+        ema("loss", decay=decay)
+
+
+@pytest.mark.parametrize("half_life", [0.0, -1.0])
+def test_ema_rejects_invalid_half_life(half_life: float) -> None:
+    with pytest.raises(ValueError):
+        ema("loss", half_life=half_life)
+
+
+def test_ema_requires_metric_key() -> None:
+    with pytest.raises(KeyError):
+        list(ema("loss", decay=0.5)([{"other": 1.0}]))
+
+
+def test_ema_rejects_non_numeric_value() -> None:
+    with pytest.raises(ValueError):
+        list(ema("loss", decay=0.5)([{"loss": "abc"}]))
+
+
+def test_ema_is_pipe_stage() -> None:
+    events = [{"loss": 10.0}, {"loss": 20.0}]
+
+    result = list(pipe(events, ema("loss", decay=0.5)))
+
+    assert "loss_ema" in result[0]
+    assert result[0]["loss"] == 10.0