[evals] Break trace label import cycle

Author: dlwh

lib/levanter/src/levanter/data/text/__init__.py

@@ -30,6 +30,7 @@
     GrugLmExample,
     LabeledLmExample,
     LOSS_IGNORE_LABEL,
+    LossLabelSpec,
     LossLabelSpan,
     grug_attention_mask_from_named,
     grug_lm_example_from_named,
@@ -88,6 +89,7 @@
     "GrugLmExample",
     "LabeledLmExample",
     "LOSS_IGNORE_LABEL",
+    "LossLabelSpec",
     "LossLabelSpan",
     "grug_attention_mask_from_named",
     "grug_lm_example_from_named",

lib/levanter/src/levanter/data/text/examples.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from dataclasses import dataclass
-from typing import Sequence
+from typing import Mapping, Sequence
 
 import jax
 import jax.numpy as jnp
@@ -35,6 +35,61 @@ class LossLabelSpan:
     label: int
 
 
+@dataclass(frozen=True)
+class LossLabelSpec:
+    """Names exclusive loss labels and defines metric rollups.
+
+    `id_to_name` names the leaf span types stored in `LabeledLmExample.loss_labels`.
+    `aggregates` maps metric names to one or more leaf label ids, so callers can
+    report both specific span types and rollups such as assistant = assistant
+    text plus assistant tool calls. If aggregates is omitted, each non-ignored
+    label id gets its own metric.
+    """
+
+    id_to_name: Mapping[int, str]
+    aggregates: Mapping[str, Sequence[int]] | None = None
+    dont_score_label: int = LOSS_IGNORE_LABEL
+
+    def __post_init__(self):
+        for label_id, name in self.id_to_name.items():
+            if not isinstance(label_id, int):
+                raise TypeError(f"label id must be an int, got {label_id!r}")
+            if not isinstance(name, str):
+                raise TypeError(f"label name for id {label_id} must be a str, got {name!r}")
+        if len(set(self.id_to_name.values())) != len(self.id_to_name):
+            raise ValueError("label names must be unique")
+
+        for name, label_ids in self._aggregate_mapping().items():
+            if not isinstance(name, str):
+                raise TypeError(f"aggregate name must be a str, got {name!r}")
+            if not label_ids:
+                raise ValueError(f"aggregate {name!r} must include at least one label id")
+            if self.dont_score_label in label_ids:
+                raise ValueError(f"aggregate {name!r} includes dont_score_label={self.dont_score_label}")
+            for label_id in label_ids:
+                if not isinstance(label_id, int):
+                    raise TypeError(f"aggregate {name!r} label id must be an int, got {label_id!r}")
+                if label_id not in self.id_to_name:
+                    raise ValueError(f"aggregate {name!r} references unknown label id {label_id}")
+
+    def _aggregate_mapping(self) -> Mapping[str, Sequence[int]]:
+        if self.aggregates is not None:
+            return self.aggregates
+        return {
+            label_name: (label_id,)
+            for label_id, label_name in self.id_to_name.items()
+            if label_id != self.dont_score_label
+        }
+
+    @property
+    def aggregate_names(self) -> tuple[str, ...]:
+        return tuple(self._aggregate_mapping().keys())
+
+    @property
+    def aggregate_label_ids(self) -> tuple[tuple[int, ...], ...]:
+        return tuple(tuple(label_ids) for label_ids in self._aggregate_mapping().values())
+
+
 def loss_labels_from_spans(
     seq_len: int,
     spans: Sequence[LossLabelSpan],

lib/levanter/src/levanter/data/text/trace_chat.py

@@ -10,7 +10,7 @@
 import numpy as np
 
 from levanter.data._preprocessor import BatchProcessor
-from levanter.eval import LossLabelSpec
+from levanter.data.text.examples import LossLabelSpec
 from levanter.tokenizers import MarinTokenizer
 
 

lib/levanter/src/levanter/eval.py

@@ -27,9 +27,9 @@
 from levanter.callbacks import StepInfo
 from levanter.data import AsyncDataset, DataLoader
 from levanter.data.text.examples import (
-    LOSS_IGNORE_LABEL,
     GrugLmExample,
     LabeledLmExample,
+    LossLabelSpec,
     named_lm_example_from_grug,
     named_lm_example_from_labeled,
 )
@@ -55,61 +55,6 @@
 BatchedTagArray = Int[Array, "... tag"]
 
 
-@dataclasses.dataclass(frozen=True)
-class LossLabelSpec:
-    """Names exclusive loss labels and defines metric rollups.
-
-    `id_to_name` names the leaf span types stored in `LabeledLmExample.loss_labels`.
-    `aggregates` maps metric names to one or more leaf label ids, so callers can
-    report both specific span types and rollups such as assistant = assistant
-    text plus assistant tool calls. If aggregates is omitted, each non-ignored
-    label id gets its own metric.
-    """
-
-    id_to_name: Mapping[int, str]
-    aggregates: Mapping[str, Sequence[int]] | None = None
-    dont_score_label: int = LOSS_IGNORE_LABEL
-
-    def __post_init__(self):
-        for label_id, name in self.id_to_name.items():
-            if not isinstance(label_id, int):
-                raise TypeError(f"label id must be an int, got {label_id!r}")
-            if not isinstance(name, str):
-                raise TypeError(f"label name for id {label_id} must be a str, got {name!r}")
-        if len(set(self.id_to_name.values())) != len(self.id_to_name):
-            raise ValueError("label names must be unique")
-
-        for name, label_ids in self._aggregate_mapping().items():
-            if not isinstance(name, str):
-                raise TypeError(f"aggregate name must be a str, got {name!r}")
-            if not label_ids:
-                raise ValueError(f"aggregate {name!r} must include at least one label id")
-            if self.dont_score_label in label_ids:
-                raise ValueError(f"aggregate {name!r} includes dont_score_label={self.dont_score_label}")
-            for label_id in label_ids:
-                if not isinstance(label_id, int):
-                    raise TypeError(f"aggregate {name!r} label id must be an int, got {label_id!r}")
-                if label_id not in self.id_to_name:
-                    raise ValueError(f"aggregate {name!r} references unknown label id {label_id}")
-
-    def _aggregate_mapping(self) -> Mapping[str, Sequence[int]]:
-        if self.aggregates is not None:
-            return self.aggregates
-        return {
-            label_name: (label_id,)
-            for label_id, label_name in self.id_to_name.items()
-            if label_id != self.dont_score_label
-        }
-
-    @property
-    def aggregate_names(self) -> tuple[str, ...]:
-        return tuple(self._aggregate_mapping().keys())
-
-    @property
-    def aggregate_label_ids(self) -> tuple[tuple[int, ...], ...]:
-        return tuple(tuple(label_ids) for label_ids in self._aggregate_mapping().values())
-
-
 @dataclasses.dataclass
 class EvalResult:
     micro_avg_loss: float  # per token across all datasets