fix: Handle missing prompts in math HF data processor and add regression test (#1219)

Signed-off-by: Zhaopeng Qiu <qiuzhaopeng@foxmail.com>

Author: zpqiu

nemo_rl/data/processors.py

@@ -108,9 +108,12 @@ def math_hf_data_processor(
     extra_env_info = {"ground_truth": user_message[1]["content"]}
 
     message_log: LLMMessageLogType = []
+    formatted_content = (
+        task_data_spec.prompt.format(problem) if task_data_spec.prompt else problem
+    )
     user_message = {
         "role": "user",
-        "content": task_data_spec.prompt.format(problem),
+        "content": formatted_content,
     }
     message: list[str] = tokenizer.apply_chat_template(  # type: ignore
         [user_message],

tests/unit/data/test_data_processor.py

@@ -18,6 +18,7 @@
 from collections import defaultdict
 
 import pytest
+import torch
 from datasets import Dataset
 
 abspath = os.path.abspath(__file__)
@@ -40,6 +41,30 @@
 from nemo_rl.models.policy import TokenizerConfig
 
 
+class DummyTokenizer:
+    def apply_chat_template(
+        self,
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        add_special_tokens=False,
+    ):
+        content = "".join(
+            f"{m.get('role', 'user')}: {m['content']}\n" for m in messages
+        )
+        if add_generation_prompt:
+            content += "assistant:"
+        return content
+
+    def __call__(self, text, return_tensors=None, add_special_tokens=False):
+        if isinstance(text, list):
+            text = "".join(text)
+        encoded = list(range(len(text)))
+        if return_tensors == "pt":
+            return {"input_ids": torch.tensor([encoded], dtype=torch.long)}
+        return {"input_ids": encoded}
+
+
 def test_math_data_processor():
     raw_dataset = Dataset.from_list(
         [
@@ -131,6 +156,37 @@ def test_math_hf_data_processor(tokenizer_name, dataset_cls):
     assert len(first_item["message_log"]) > 0
 
 
+def test_math_hf_data_processor_without_prompt():
+    datum_dict = {
+        "messages": [
+            {"role": "user", "content": "Solve 1+1."},
+            {"role": "assistant", "content": "2"},
+        ],
+        "task_name": "math",
+    }
+    tokenizer = DummyTokenizer()
+
+    math_task_spec = TaskDataSpec(
+        task_name="math",
+        prompt_file=None,
+        system_prompt_file=None,
+    )
+
+    result = math_hf_data_processor(
+        datum_dict=datum_dict,
+        task_data_spec=math_task_spec,
+        tokenizer=tokenizer,
+        max_seq_length=128,
+        idx=0,
+    )
+
+    assert result["extra_env_info"]["ground_truth"] == "2"
+    assert result["loss_multiplier"] == 1.0
+    assert len(result["message_log"]) == 1
+    assert result["message_log"][0]["role"] == "user"
+    assert "Solve 1+1." in result["message_log"][0]["content"]
+
+
 @pytest.fixture
 def system_prompt_file(request):
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as file: