Repository: oumi-ai/oumi
Pull request: add debug logging capabilities to collators
# Description

<!--
Thank you for contributing to Oumi! Before sending your PR out for review, please take a quick read through this template.

When your PR is merged, its title will appear in our release notes. Make sure your title gives a clear description of your change!

After you've updated your title, please replace this section with a detailed description of your change. Include as much context as possible so your reviewers can easily understand *what* you're changing and *why*.
The more information you provide, the faster we can review your change!
-->
<!--↓↓↓↓↓↓↓↓↓↓ Describe your change below ↓↓↓↓↓↓↓↓↓↓-->
This PR adds a comprehensive debug-logging feature to Oumi’s data collators, giving developers clear, step-by-step visibility into how raw text becomes model-ready tokens. By setting `debug=True`, you’ll get both console output and an HTML report (organized by session) that shows:

- **Raw Examples**: Original dataset records before any processing  
- **Formatted Text**: After applying prompts or templates  
- **Token Breakdown**: Each token’s ID, text, and position  
- **Collated Inputs**: Final batched tensors passed to the model  

## Key Changes

1. **Debug Flag**  
   - Added a `debug: bool` parameter to all collator builders and classes.  

2. **`debug_utils.py`**  
   - New utility module with reusable logging functions.  
   - `log_example_for_debugging()` records:
     - **Raw example** (pre-tokenization)
     - **Post-format example** (after any prompt/template)
     - **Token list** (IDs, decoded strings, special tokens)
     - **Final model inputs** (padded & batched tensors)

3. **Console & HTML Output**  
   - **Console**: Uses the existing `logger.debug()` interface.  
   - **HTML Reports**:
     - Written to `debug_logs/<session_id>/` with a timestamp.
     - A `latest.html` symlink always points to the newest report.
     - Clean, sectioned layouts with preformatted JSON for easy browser inspection.

4. **Updated Collators**  
   - `TextCollatorWithPadding`  
   - `TextCompletionsCollatorWithPadding`  
   Both now accept `debug=True` and integrate with the new utilities.

## Usage

```python
from oumi.data_collators import TextCollatorWithPadding

collator = TextCollatorWithPadding(
    tokenizer=my_tokenizer,
    max_length=512,
    debug=True  # ← enables detailed logging & HTML report
)
```
<!--↑↑↑↑↑↑↑↑↑↑ Describe your change above ↑↑↑↑↑↑↑↑↑↑-->

## Related issues

<!--
Make sure to list any relevant related issues to your change. More often than not this will be the single issue fixed by your PR.
-->
<!--↓↓↓↓↓↓↓↓↓↓ List your related issues below ↓↓↓↓↓↓↓↓↓↓-->

Fixes #1369

<!--↑↑↑↑↑↑↑↑↑↑ List your related issues above ↑↑↑↑↑↑↑↑↑↑-->

## Before submitting

- [ ] This PR only changes documentation. (You can ignore the following checks in that case)
- [X] Did you read the [contributor guideline](https://github.com/oumi-ai/oumi/blob/main/CONTRIBUTING.md) Pull Request guidelines?
- [X] Did you link the issue(s) related to this PR in the section above?
- [x] Did you add / update tests where needed?

## Reviewers

At least one review from a member of `oumi-ai/oumi-staff` is required.

<!-- Add `oumi-ai/oumi-staff` as a reviewer when your PR is ready for review.

You are also welcome to add individual members of `oumi-ai/oumi-staff` as reviewers.

If no one has reviewed your PR after several days, feel free to add a comment tagging specific reviewers.

 -->


Problem statement:
[Feature] Output tokenized example for debugging during training
### Feature request

When running training, it'd be useful to log a single example from the dataset that has been tokenized for debugging.

For example:
1. Take the first example or batch in the training set (peek the top of the dataloader?)
2. Print all the following:
a. Raw Example
b. Formatted Example
c. Tokenized Example (list of tuples of <token_id, decoded_token>)
d. Model Input (tokens with attention mask and labels)

Ideally this should either be a debug log or its output explicitly controlled by a flag.

Likely, all parts except the model input can be done when loading the dataset. The model input will need to be done inside the collator(s).

### Motivation / references

Concrete example using `HuggingFaceTB/SmolLM2-1.7B-Instruct` tokenizer, with model setup to be trained on **completions only**:

Raw Example
```
[
  {'role': 'system', 'content': 'This is a system instruction.'},
  {'role': 'user', 'content': 'This is a user instruction'},
  {'role': 'assistant', 'content': 'This is an assistant response.'},
]
```

Formatted Example
```

system
This is a system instruction.


user
This is a user instruction


assistant
This is an assistant response.

```

Tokenized Example
```
[(1, '
'), (9690, 'system'), (198, '\n'), (1348, 'This'), (314, ' is'), (253, ' a'), (817, ' system'), (5785, ' instruction'), (30, '.'), (2, '
'), (198, '\n'), (1, '
'), (4093, 'user'), (198, '\n'), (1348, 'This'), (314, ' is'), (253, ' a'), (2914, ' user'), (5785, ' instruction'), (2, '
'), (198, '\n'), (1, '
'), (520, 'ass'), (9531, 'istant'), (198, '\n'), (1348, 'This'), (314, ' is'), (354, ' an'), (11173, ' assistant'), (2426, ' response'), (30, '.'), (2, '
'), (198, '\n')]
```

Model Input Example
```
{
  'input': [1, 9690, 198, 1348, 314, 253, 817, 5785, 30, 2, 198, 1, 4093, 198, 1348, 314, 253, 2914, 5785, 2, 198, 1, 520, 9531, 198, 1348, 314, 354, 11173, 2426, 30, 2, 198],
  'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
  'labels': [-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 1348, 314, 354, 11173, 2426, 30, 2, 198],
}
```

### Your contribution

Example code:
```
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained('HuggingFaceTB/SmolLM2-1.7B-Instruct')

chat = [
  {'role': 'system', 'content': 'This is a system instruction.'},
  {'role': 'user', 'content': 'This is a user instruction'},
  {'role': 'assistant', 'content': 'This is an assistant response.'},
]

formatted = tokenizer.apply_chat_template(chat, tokenize=False)
print(formatted)

tokens = tokenizer.apply_chat_template(chat)
decoded = [tokenizer.decode(t) for t in tokens]
print(list(zip(tokens, decoded)))
```

Hints from issue discussion:
in addition to text output (`print()`) , it may also be useful to write the information into e.g., a nicely-formatted HTML file 
Hi @nikg4, I would like to work on this issue if it is still active!
@aniruddh-alt  of course!

Commit message: Merge branch 'main' into fix/1369-debug-token-logging

Patch to review:
diff --git a/src/oumi/builders/collators.py b/src/oumi/builders/collators.py
index 592951a3c..46b3a92bc 100644
--- a/src/oumi/builders/collators.py
+++ b/src/oumi/builders/collators.py
@@ -40,6 +40,7 @@ def build_data_collator(
     *,
     max_length: Optional[int],
     label_ignore_index: Optional[int] = constants.LABEL_IGNORE_INDEX,
+    debug: bool = False,
     **kwargs,
 ) -> Callable:
     """Builds a data collator based on the given collator name.
@@ -62,6 +63,7 @@ def build_data_collator(
             PyTorch convention is to use -100 as the `ignore_index` label. Refer to
             the `ignore_index` parameter of `torch.nn.CrossEntropyLoss()`
             for more details.
+        debug: If True, logs a single example for debugging purposes.
         **kwargs: Additional keyword arguments to pass to the collator constructor.
 
     Returns:
@@ -91,16 +93,18 @@ def build_data_collator(
                 + f" tokenizer's model maximum length ({tokenizer.model_max_length})"
             )
 
+    collator = None
     if collator_name == "text_with_padding":
-        return TextCollatorWithPadding(
+        collator = TextCollatorWithPadding(
             tokenizer=tokenizer,
             max_length=max_length,
             truncation=enable_truncation,
             label_ignore_index=label_ignore_index,
+            debug=debug,
             **kwargs,
         )
     elif collator_name == "vision_language_with_padding":
-        return VisionLanguageCollatorWithPadding(
+        collator = VisionLanguageCollatorWithPadding(
             tokenizer=tokenizer,
             max_length=max_length,
             truncation=enable_truncation,
@@ -112,7 +116,7 @@ def build_data_collator(
         if not processor_name:
             raise ValueError(f"Empty processor_name for '{collator_name}'")
         processor_kwargs = kwargs.pop("processor_kwargs", None)
-        return VisionLanguageSftCollator(
+        collator = VisionLanguageSftCollator(
             tokenizer=tokenizer,
             processor_name=processor_name,
             processor_kwargs=processor_kwargs,
@@ -122,17 +126,20 @@ def build_data_collator(
             **kwargs,
         )
     elif collator_name == "text_completions_only_with_padding":
-        return TextCompletionsCollatorWithPadding(
+        collator = TextCompletionsCollatorWithPadding(
             tokenizer=tokenizer,
             instruction_prefix="
user
\n\n",
             response_prefix="
assistant
\n\n",
+            debug=debug,
         )
 
-    raise ValueError(f"Unknown data collator name: '{collator_name}'")
+    if collator is None:
+        raise ValueError(f"Unknown data collator name: '{collator_name}'")
+    return collator
 
 
 def build_collator_from_config(
-    config: TrainingConfig, tokenizer: Optional[BaseTokenizer]
+    config: TrainingConfig, tokenizer: Optional[BaseTokenizer], debug: bool = False
 ) -> Optional[Callable]:
     """Creates data collator if specified in config."""
     train_split = config.data.get_split(DatasetSplit.TRAIN)
@@ -190,5 +197,6 @@ def build_collator_from_config(
         tokenizer=tokenizer,
         max_length=config.model.model_max_length,
         label_ignore_index=label_ignore_index,
+        debug=debug,
         **collator_kwargs,
     )
diff --git a/src/oumi/core/collators/text_collator_with_padding.py b/src/oumi/core/collators/text_collator_with_padding.py
index 671dac53b..6648699ca 100644
--- a/src/oumi/core/collators/text_collator_with_padding.py
+++ b/src/oumi/core/collators/text_collator_with_padding.py
@@ -16,6 +16,7 @@
 from typing import Any, NamedTuple, Optional
 
 from oumi.core.tokenizers.base_tokenizer import BaseTokenizer
+from oumi.utils.debug_utils import log_example_for_debugging
 from oumi.utils.logging import logger
 from oumi.utils.torch_utils import (
     create_ones_like,
@@ -50,6 +51,7 @@ def __init__(
         truncation: bool = False,
         label_ignore_index: Optional[int] = None,
         max_variable_sized_dims: int = 1,
+        debug: bool = False,
     ):
         """Custom collator for text LLM training.
 
@@ -65,6 +67,7 @@ def __init__(
             Normally, it's 1 (sequence length dimension), but can sometimes be higher
             e.g., 2 for "cross_attention_mask" for VLM-s with multi-image inputs.
             Negative value mean `Unlimited`.
+        debug: Whether to log a debug example.
         """
         self._max_length: Optional[int] = (
             int(max_length) if max_length is not None and max_length > 0 else None
@@ -91,6 +94,10 @@ def __init__(
         self._max_input_ids_length: int = 0
         self._max_previously_logged_input_ids_length: int = 0
         self._max_variable_sized_dims: int = max_variable_sized_dims
+        self._debug: bool = debug
+        # Track if we've already logged an example
+        self._has_logged_example: bool = False
+        self._tokenizer = tokenizer  # Store tokenizer for debugging
 
     def _collate_simple(
         self,
@@ -223,6 +230,51 @@ def __call__(self, batch) -> dict[str, Any]:
         if labels_on:
             combined_batch[_LABELS_KEY] = collated_text_inputs[_LABELS_KEY]
 
+        # If debug is on and we haven't logged an example yet, log the first example
+        if self._debug and not self._has_logged_example and len(batch) > 0:
+            first_input_ids = combined_batch[_INPUT_IDS_KEY][0]
+            formatted_example = self._tokenizer.decode(
+                first_input_ids, skip_special_tokens=False
+            )
+
+            tokenized_example = [
+                (
+                    int(tid.item() if hasattr(tid, "item") else tid),
+                    self._tokenizer.decode([tid])
+                    if hasattr(tid, "item")
+                    else self._tokenizer.decode(tid),
+                )
+                for tid in first_input_ids
+            ]
+
+            model_input = {
+                "input_ids": (
+                    first_input_ids.tolist()
+                    if hasattr(first_input_ids, "tolist")
+                    else first_input_ids
+                ),
+                "attention_mask": (
+                    combined_batch[_ATTENTION_MASK_KEY][0].tolist()
+                    if hasattr(combined_batch[_ATTENTION_MASK_KEY][0], "tolist")
+                    else combined_batch[_ATTENTION_MASK_KEY][0]
+                ),
+            }
+
+            if labels_on:
+                model_input["labels"] = (
+                    combined_batch[_LABELS_KEY][0].tolist()
+                    if hasattr(combined_batch[_LABELS_KEY][0], "tolist")
+                    else combined_batch[_LABELS_KEY][0]
+                )
+
+            # Mark that we've logged an example to avoid logging again
+            self._has_logged_example = True
+            log_example_for_debugging(
+                raw_example=batch[0],
+                formatted_example=str(formatted_example),
+                tokenized_example=tokenized_example,
+                model_input=model_input,
+            )
         return combined_batch
 
     def _update_max_lengths_and_log(self, *, max_input_ids_length: int):
diff --git a/src/oumi/core/collators/text_completions_collator_with_padding.py b/src/oumi/core/collators/text_completions_collator_with_padding.py
index f489a084c..b56e16f8b 100644
--- a/src/oumi/core/collators/text_completions_collator_with_padding.py
+++ b/src/oumi/core/collators/text_completions_collator_with_padding.py
@@ -17,13 +17,18 @@
 import trl
 
 from oumi.core.tokenizers.base_tokenizer import BaseTokenizer
+from oumi.utils.debug_utils import log_example_for_debugging
 
 _INPUT_IDS_KEY = "input_ids"
 
 
 class TextCompletionsCollatorWithPadding:
     def __init__(
-        self, tokenizer: BaseTokenizer, instruction_prefix: str, response_prefix: str
+        self,
+        tokenizer: BaseTokenizer,
+        instruction_prefix: str,
+        response_prefix: str,
+        debug: bool = False,
     ):
         """Custom collator for text LLM training.
 
@@ -31,6 +36,7 @@ def __init__(
         tokenizer: The tokenizer used for encoding the data.
         instruction_prefix: The prefix marking the beginning of the user instruction.
         response_prefix: The prefix marking the beginning of the assistant response.
+        debug: If True, enables debug mode for logging.
         """
         self._default_collator = trl.DataCollatorForCompletionOnlyLM(
             tokenizer=tokenizer,
@@ -41,6 +47,9 @@ def __init__(
         if not hasattr(tokenizer, "pad_token_id") or tokenizer.pad_token_id is None:
             raise RuntimeError("Tokenizer doesn't define `pad_token_id`.")
 
+        self._debug = debug
+        self._has_logged_example = False
+
     def _collate(self, inputs: list[Any]) -> dict[str, Any]:
         result = self._default_collator(inputs)
         return result
@@ -64,4 +73,30 @@ def __call__(self, batch) -> dict[str, Any]:
         # Collate batch prompts.
         collated_text_inputs = self._collate(batch)
 
+        if self._debug and not self._has_logged_example:
+            # Log the first example for debugging
+            raw_example = batch[0]
+
+            # Get the formatted text from the tokenizer's encoding
+            formatted_example = self._default_collator.tokenizer.decode(
+                raw_example[_INPUT_IDS_KEY], skip_special_tokens=False
+            )
+
+            # Tokenize the formatted example
+            tokenized_ids = raw_example[_INPUT_IDS_KEY]
+            # Create tokenized example pairs
+            tokenized_example = [
+                (token_id, self._default_collator.tokenizer.decode([token_id]))
+                for token_id in tokenized_ids
+            ]
+
+            # Get model input (same as collated_text_inputs but for a single example)
+            model_input = self._collate([raw_example])
+
+            # Log all components for debugging
+            log_example_for_debugging(
+                raw_example, formatted_example, tokenized_example, model_input
+            )
+            self._has_logged_example = True
+
         return collated_text_inputs
diff --git a/src/oumi/utils/debug_utils.py b/src/oumi/utils/debug_utils.py
new file mode 100644
index 000000000..b498547bd
--- /dev/null
+++ b/src/oumi/utils/debug_utils.py
@@ -0,0 +1,95 @@
+import datetime
+import json
+import uuid
+from pathlib import Path
+from typing import Any
+
+from oumi.utils.logging import logger
+
+
+def log_example_for_debugging(
+    raw_example: Any,
+    formatted_example: str,
+    tokenized_example: list[tuple[int, str]],
+    model_input: dict[str, Any],
+) -> None:
+    """Logs an example of the data in each step for debugging purposes.
+
+    Args:
+        raw_example: The raw example from the dataset.
+        formatted_example: The formatted example after processing.
+        tokenized_example: The tokenized example after tokenization.
+        model_input: The final model input after collating.
+    """
+    # Log to debug file
+    logger.debug("Raw example: %s", raw_example)
+    logger.debug("Formatted example: %s", formatted_example)
+    logger.debug("Tokenized example: %s", tokenized_example)
+    logger.debug("Model input: %s", model_input)
+
+    # Generate timestamp for the debug file
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    session_id = str(uuid.uuid4())[:6]
+
+    # Format data for HTML display
+    def format_for_html(obj):
+        try:
+            return json.dumps(obj, indent=2, default=str)
+        except Exception:
+            return str(obj)
+
+    # Create simplified HTML content
+    html_content = f"""<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="UTF-8">
+    <title>Debug Log {timestamp}</title>
+    <style>
+        body {{ font-family: sans-serif; margin: 20px; }}
+        h1 {{ color: #333; }}
+        .info {{ color: #666; margin-bottom: 20px; }}
+        .section {{ margin-bottom: 30px; }}
+        h2 {{ background: #eee; padding: 8px; }}
+        pre {{ background: #f8f8f8; padding: 10px; overflow: auto;
+             border: 1px solid #ddd; }}
+        .copy-btn {{ float: right; padding: 3px 8px; background: #eee;
+                  border: 1px solid #ccc; cursor: pointer; }}
+    </style>
+</head>
+<body>
+    <h1>Oumi Debug Information</h1>
+    <div class="section">
+        <h2>Raw Example</h2>
+        <pre id="raw">{format_for_html(raw_example)}</pre>
+    </div>
+
+    <div class="section">
+        <h2>Formatted Example</h2>
+        <pre id="formatted">{formatted_example}</pre>
+    </div>
+
+    <div class="section">
+        <h2>Tokenized Example</h2>
+        <pre id="tokenized">{format_for_html(tokenized_example)}</pre>
+    </div>
+
+    <div class="section">
+        <h2>Model Input</h2>
+        <pre id="model">{format_for_html(model_input)}</pre>
+    </div>
+</body>
+</html>"""
+
+    # Create output directory and write HTML files
+    output_dir = "debug_logs"
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+    # Write to timestamped file
+    output_file = Path(output_dir) / f"debug_logs_{timestamp}_{session_id}.html"
+    with open(output_file, "w") as f:
+        f.write(html_content)
+
+    # Also update the latest.html file
+    latest_file = Path(output_dir) / "latest.html"
+    with open(latest_file, "w") as f:
+        f.write(html_content)


Code review comments:
diff_hunk: @@ -0,0 +1,95 @@
+import datetime
+import json
+import uuid
+from pathlib import Path
+from typing import Any
+
+from oumi.utils.logging import logger
+
+
+def log_example_for_debugging(
+    raw_example: Any,
+    formatted_example: str,
+    tokenized_example: list[tuple[int, str]],
+    model_input: dict[str, Any],
+) -> None:
+    """Logs an example of the data in each step for debugging purposes.
+
+    Args:
+        raw_example: The raw example from the dataset.
+        formatted_example: The formatted example after processing.
+        tokenized_example: The tokenized example after tokenization.
+        model_input: The final model input after collating.
+    """
+    # Log to debug file
+    logger.debug("Raw example: %s", raw_example)
+    logger.debug("Formatted example: %s", formatted_example)
+    logger.debug("Tokenized example: %s", tokenized_example)
+    logger.debug("Model input: %s", model_input)
+
+    # Generate timestamp for the debug file
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    session_id = str(uuid.uuid4())[:6]
+
+    # Format data for HTML display
+    def format_for_html(obj):
+        try:
+            return json.dumps(obj, indent=2, default=str)
+        except Exception:
+            return str(obj)
path: src/oumi/utils/debug_utils.py
text: @user1:
Saving as HTML is likely overkill here. I think calling `logger.debug` is likely sufficient.

@author:
Thank you for the comment! I will remove the HTML part from the function.

diff_hunk: @@ -91,16 +93,18 @@ def build_data_collator(
                 + f" tokenizer's model maximum length ({tokenizer.model_max_length})"
             )
 
+    collator = None
     if collator_name == "text_with_padding":
-        return TextCollatorWithPadding(
+        collator = TextCollatorWithPadding(
path: src/oumi/builders/collators.py
text: @user1:
Are the changes to these lines necessary? If not, let's leave them as is. Early returns are less error prone than a single None check. 

@author:
I have reverted back to early returns

diff_hunk: @@ -64,4 +73,30 @@ def __call__(self, batch) -> dict[str, Any]:
         # Collate batch prompts.
         collated_text_inputs = self._collate(batch)
 
+        if self._debug and not self._has_logged_example:
+            # Log the first example for debugging
+            raw_example = batch[0]
path: src/oumi/core/collators/text_completions_collator_with_padding.py
text: @user1:
Can this new code block be moved into a helper function? it's relatively complex/long

add unit tests for new functionality

@author:
Thank you for the comment, I will create a helper function and add unit tests as well.

diff_hunk: @@ -223,6 +230,51 @@ def __call__(self, batch) -> dict[str, Any]:
         if labels_on:
             combined_batch[_LABELS_KEY] = collated_text_inputs[_LABELS_KEY]
 
+        # If debug is on and we haven't logged an example yet, log the first example
+        if self._debug and not self._has_logged_example and len(batch) > 0:
+            first_input_ids = combined_batch[_INPUT_IDS_KEY][0]
path: src/oumi/core/collators/text_collator_with_padding.py
text: @user1:
Can this new code block be moved into a helper function? it's relatively complex/long

+ add unit tests for new functionality 

@author:
I have added helper function and unit tests