[evals] Fix trace chat template rendering

Author: dlwh

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

@@ -422,7 +422,7 @@ def metadata(self) -> dict[str, Any]:
             "include_final_assistant_tag": self.include_final_assistant_tag,
             "parse_text_tool_calls": self.parse_text_tool_calls,
             "label_spec": {
-                "id_to_name": self.label_spec.id_to_name,
+                "id_to_name": {str(label_id): name for label_id, name in self.label_spec.id_to_name.items()},
                 "aggregates": {name: list(label_ids) for name, label_ids in self.label_spec.aggregates.items()},
             },
         }

lib/levanter/src/levanter/tokenizers.py

@@ -195,12 +195,18 @@ def _passthrough(caller: jinja2.runtime.Macro) -> str:
 
 def _make_jinja_env(extensions: list[type]) -> jinja2.Environment:
     """Create a jinja2 environment matching HF's template rendering settings."""
-    return jinja2.Environment(
+    env = jinja2.Environment(
         undefined=jinja2.StrictUndefined,
         trim_blocks=True,
         lstrip_blocks=True,
         extensions=extensions,
     )
+    env.globals["raise_exception"] = _raise_chat_template_exception
+    return env
+
+
+def _raise_chat_template_exception(message: str) -> None:
+    raise jinja2.exceptions.TemplateError(message)
 
 
 def _chat_template_block_contents(block: str) -> str:

lib/levanter/tests/test_tokenizers.py

@@ -1103,6 +1103,28 @@ def test_apply_chat_template_with_masks_batch(backend_tokenizer):
     assert len(result["assistant_masks"]) == 2
 
 
+@requires_model
+def test_apply_chat_template_with_masks_supports_raise_exception(backend_tokenizer):
+    template = """\
+{%- for message in messages -%}
+{%- if message.role == 'tool' -%}
+{{- raise_exception('tool messages are not supported') -}}
+{%- endif -%}
+{{ message.content }}
+{%- endfor -%}
+"""
+    backend_tokenizer.apply_chat_template_with_masks(
+        [[{"role": "user", "content": "hello"}]],
+        chat_template=template,
+    )
+
+    with pytest.raises(Exception, match="tool messages are not supported"):
+        backend_tokenizer.apply_chat_template_with_masks(
+            [[{"role": "tool", "content": "result"}]],
+            chat_template=template,
+        )
+
+
 # ---------------------------------------------------------------------------
 # 10. Code samples and real-world text
 # ---------------------------------------------------------------------------