docs and naming

Author: Felipe Mello

src/forge/data/__init__.py

@@ -4,6 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 from .collate import collate_packed
+from .metric_transform import DefaultDatasetMetricTransform, MetricTransform
 from .utils import CROSS_ENTROPY_IGNORE_IDX
 
-__all__ = ["collate_packed", "CROSS_ENTROPY_IGNORE_IDX"]
+__all__ = [
+    "collate_packed",
+    "CROSS_ENTROPY_IGNORE_IDX",
+    "MetricTransform",
+    "DefaultDatasetMetricTransform",
+]

src/forge/data/dataset_metrics/__init__.py

@@ -12,7 +12,7 @@
 from datasets import load_dataset
 from datasets.distributed import split_dataset_by_node
 
-from forge.data.dataset_metrics import DefaultTrainingMetricTransform, MetricTransform
+from forge.data.metric_transform import DefaultDatasetMetricTransform, MetricTransform
 from forge.interfaces import Transform
 from forge.observability.metrics import Metric, Reduce
 
@@ -82,7 +82,7 @@ def __init__(
         self._weight = weight if weight is not None else 1.0
 
         # Create default transform if not provided
-        self._metric_transform = metric_transform or DefaultTrainingMetricTransform()
+        self._metric_transform = metric_transform or DefaultDatasetMetricTransform()
 
         # Auto-generate dataset name if not provided
         if dataset_name is None:

src/forge/data/dataset_metrics/metric_transform.py

@@ -9,7 +9,7 @@
 import torch
 
 from forge.data import CROSS_ENTROPY_IGNORE_IDX
-from forge.data.dataset_metrics import DefaultTrainingMetricTransform
+from forge.data.metric_transform import DefaultDatasetMetricTransform
 from forge.data.utils import mask_messages, TuneMessage
 from forge.interfaces import Transform
 
@@ -200,7 +200,7 @@ def sft_iterable_dataset(
         message_transform=message_transform,
         model_transform=model_transform,
         output_transform=output_transform,
-        metric_transform=DefaultTrainingMetricTransform(),
+        metric_transform=DefaultDatasetMetricTransform(),
         shuffle_buffer_size=shuffle_buffer_size,
         weight=weight,
         seed=seed,

src/forge/data/datasets/hf_dataset.py

@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any
+
+from forge.interfaces import Transform
+from forge.observability.metrics import Metric, Reduce
+
+
+class MetricTransform(Transform):
+    """
+    Base class for transforms that collect observability metrics from dataset samples.
+
+    This class provides a foundation for implementing dataset-level metric collection
+    during data processing pipelines. Subclasses should override the __call__ method
+    to add specific metrics to each sample that passes through the transform.
+
+    Metrics are collected as `forge.observability.metrics.Metric` objects and made available
+    in batch["metrics"].
+
+    Attributes:
+        source (str, optional): The source name for metrics, typically the dataset name.
+            This is used as a prefix in metric keys to distinguish metrics from different
+            data sources.
+
+    Example:
+        >>> transform = SomeMetricTransform()
+        >>> transform.set_source("training_data")
+        >>> processed_sample = transform(sample)
+        >>> # Metrics are automatically added to sample["metrics"]
+    """
+
+    def __init__(self):
+        self.source = None
+
+    def set_source(self, source: str):
+        """Set the source name for metrics (typically the dataset name)."""
+        self.source = source
+
+    def __call__(self, sample: dict[str, Any]) -> dict[str, Any]:
+        """Transform a sample by adding metrics to it."""
+        return sample
+
+
+class DefaultDatasetMetricTransform(MetricTransform):
+    """
+    Collects basic dataset processing metrics during data pipeline execution.
+
+    Metrics collected:
+    - samples_processed: Total number of samples that have passed through this transform (SUM)
+    - tokens_processed: Total number of tokens processed across all samples (SUM)
+    - mean_seq_len: Average sequence length across samples (MEAN)
+    - max_seq_len: Maximum sequence length observed (MAX)
+    - min_seq_len: Minimum sequence length observed (MIN)
+
+    Note: Token-related metrics are only collected if the sample contains a 'tokens' field.
+    Sequence length is measured as the number of tokens in each sample.
+
+    Example:
+        >>> collector = DefaultDatasetMetricTransform()
+        >>> collector.set_source("training_data")
+        >>> sample = {"tokens": ["hello", "world"]}
+        >>> processed_sample = collector(sample)
+        >>> # Metrics are automatically added to processed_sample["metrics"]
+    """
+
+    def __call__(self, sample: dict[str, Any]) -> dict[str, Any]:
+        if "metrics" not in sample:
+            sample["metrics"] = []
+
+        source_name = self.source or "dataset"
+
+        # Add samples_processed metric
+        sample["metrics"].append(
+            Metric(
+                key=f"dataset/{source_name}/samples_processed",
+                value=1,
+                reduction=Reduce.SUM,
+            )
+        )
+
+        # Add token-based metrics if tokens are present
+        if "tokens" in sample:
+            token_count = len(sample.get("tokens", []))
+
+            sample["metrics"].extend(
+                [
+                    Metric(
+                        key=f"dataset/{source_name}/tokens_processed",
+                        value=token_count,
+                        reduction=Reduce.SUM,
+                    ),
+                    Metric(
+                        key=f"dataset/{source_name}/mean_seq_len",
+                        value=token_count,
+                        reduction=Reduce.MEAN,
+                    ),
+                    Metric(
+                        key=f"dataset/{source_name}/max_seq_len",
+                        value=token_count,
+                        reduction=Reduce.MAX,
+                    ),
+                    Metric(
+                        key=f"dataset/{source_name}/min_seq_len",
+                        value=token_count,
+                        reduction=Reduce.MIN,
+                    ),
+                ]
+            )
+
+        return sample

src/forge/data/datasets/sft_dataset.py

@@ -25,9 +25,9 @@
 
 import pytest
 import torch.distributed as dist
-from forge.data.dataset_metrics import DefaultTrainingMetricTransform
 
 from forge.data.datasets import HfIterableDataset
+from forge.data.metric_transform import DefaultDatasetMetricTransform
 from torch.testing._internal.common_fsdp import FSDPTest
 
 from torchdata.stateful_dataloader import StatefulDataLoader
@@ -93,7 +93,7 @@ def _create_dataset(
             dataset_name=dataset_name,
             seed=SEED,
             shuffle_buffer_size=10 if shuffle else 0,
-            metric_transform=DefaultTrainingMetricTransform(),
+            metric_transform=DefaultDatasetMetricTransform(),
             num_shards_per_rank=2,
             **kwargs,
         )

src/forge/data/metric_transform.py

@@ -28,9 +28,9 @@
 
 import torch
 import torch.distributed as dist
-
-from forge.data.dataset_metrics import DefaultTrainingMetricTransform
 from forge.data.datasets import HfIterableDataset, InterleavedDataset
+
+from forge.data.metric_transform import DefaultDatasetMetricTransform
 from torch.testing._internal.common_fsdp import FSDPTest
 from torchdata.stateful_dataloader import StatefulDataLoader
 
@@ -114,7 +114,7 @@ def _create_dataset(
             dataset_name=dataset_name,
             seed=SEED,
             shuffle_buffer_size=10 if shuffle else 0,
-            metric_transform=DefaultTrainingMetricTransform(),
+            metric_transform=DefaultDatasetMetricTransform(),
             num_shards_per_rank=2,
             **kwargs,
         )
@@ -308,38 +308,38 @@ def test_metrics_aggregation(
             if "metrics" in sample:
                 collected_metrics.extend(sample["metrics"])
 
-        # Count metrics by dataset name
-        ds1_samples_seen = sum(
+        # Count metrics by dataset name (using new metric key)
+        ds1_samples_processed = sum(
             1
             for m in collected_metrics
-            if hasattr(m, "key") and "dataset/ds1/samples_seen" in m.key
+            if hasattr(m, "key") and "dataset/ds1/samples_processed" in m.key
         )
-        ds2_samples_seen = sum(
+        ds2_samples_processed = sum(
             1
             for m in collected_metrics
-            if hasattr(m, "key") and "dataset/ds2/samples_seen" in m.key
+            if hasattr(m, "key") and "dataset/ds2/samples_processed" in m.key
         )
-        ds3_samples_seen = sum(
+        ds3_samples_processed = sum(
             1
             for m in collected_metrics
-            if hasattr(m, "key") and "dataset/ds3/samples_seen" in m.key
+            if hasattr(m, "key") and "dataset/ds3/samples_processed" in m.key
         )
 
         # All datasets should have contributed samples
-        assert ds1_samples_seen > 0, "ds1 should have contributed samples"
-        assert ds2_samples_seen > 0, "ds2 should have contributed samples"
-        assert ds3_samples_seen > 0, "ds3 should have contributed samples"
+        assert ds1_samples_processed > 0, "ds1 should have contributed samples"
+        assert ds2_samples_processed > 0, "ds2 should have contributed samples"
+        assert ds3_samples_processed > 0, "ds3 should have contributed samples"
 
         # Total samples should equal what we processed
         calculated_total_samples = (
-            ds1_samples_seen + ds2_samples_seen + ds3_samples_seen
+            ds1_samples_processed + ds2_samples_processed + ds3_samples_processed
         )
         assert calculated_total_samples == total_samples
 
         # Test that ratios are approximately correct based on nested weighting
-        ds1_ratio = ds1_samples_seen / total_samples
-        ds2_ratio = ds2_samples_seen / total_samples
-        ds3_ratio = ds3_samples_seen / total_samples
+        ds1_ratio = ds1_samples_processed / total_samples
+        ds2_ratio = ds2_samples_processed / total_samples
+        ds3_ratio = ds3_samples_processed / total_samples
 
         # Expected ratios based on nested weighting:
         # Inner weights: ds1=0.2, ds2=0.8 -> inner total=1.0
@@ -518,7 +518,7 @@ def create_dataset():
                     split="train",
                     dataset_name="ds1",
                     shuffle_buffer_size=0,  # No shuffle for determinism
-                    metric_transform=DefaultTrainingMetricTransform(),
+                    metric_transform=DefaultDatasetMetricTransform(),
                     num_shards_per_rank=2,
                     weight=0.3,
                 )
@@ -528,7 +528,7 @@ def create_dataset():
                     split="train",
                     dataset_name="ds2",
                     shuffle_buffer_size=0,  # No shuffle for determinism
-                    metric_transform=DefaultTrainingMetricTransform(),
+                    metric_transform=DefaultDatasetMetricTransform(),
                     num_shards_per_rank=2,
                     weight=0.7,
                 )
@@ -538,7 +538,7 @@ def create_dataset():
                     split="train",
                     dataset_name="ds3",
                     shuffle_buffer_size=0,  # No shuffle for determinism
-                    metric_transform=DefaultTrainingMetricTransform(),
+                    metric_transform=DefaultDatasetMetricTransform(),
                     num_shards_per_rank=2,
                     weight=1.0,
                 )