fix: Fix incorrect indexing of message which cuts off user message when we… (#815)

Signed-off-by: Parth Chadha <pchadha@nvidia.com>

Author: parthchadha

examples/run_grpo_math.py

@@ -90,7 +90,7 @@ def hf_data_processor(
         add_special_tokens=False,
     )
     user_message["token_ids"] = tokenizer(message, return_tensors="pt")["input_ids"][0]
-    user_message["content"] = message[0]
+    user_message["content"] = message
     message_log.append(user_message)
 
     length = sum(len(m["token_ids"]) for m in message_log)

tests/unit/utils/test_logger.py

@@ -26,6 +26,7 @@
     TensorboardLogger,
     WandbLogger,
     flatten_dict,
+    print_message_log_samples,
 )
 
 
@@ -1441,3 +1442,23 @@ def test_log_hyperparams_with_mlflow(
         mock_wandb_instance.log_hyperparams.assert_called_once_with(params)
         mock_tb_instance.log_hyperparams.assert_called_once_with(params)
         mock_mlflow_instance.log_hyperparams.assert_called_once_with(params)
+
+
+def test_print_message_log_samples(capsys):
+    """Test that print_message_log_samples displays full content correctly."""
+    # Test message with full content (verifies our bug fix)
+    message_logs = [
+        [
+            {"role": "user", "content": "What is 2+2?"},
+            {"role": "assistant", "content": "2+2 = 4"},
+        ]
+    ]
+    rewards = [1.0]
+
+    print_message_log_samples(message_logs, rewards, num_samples=1, step=0)
+
+    captured = capsys.readouterr()
+    # Verify content is displayed properly
+    assert "What is 2+2?" in captured.out
+    assert "2+2 = 4" in captured.out
+    assert "Sample 1 | Reward: 1.0000" in captured.out