Fix TRL formatter (#378)

* Fix TRL formatter

* No need for seperate return

* Improve safety

Author: lukehinds

deepfabric/formatters/builtin/trl_sft_tools.py

@@ -37,6 +37,7 @@
 - https://www.stephendiehl.com/posts/fine_tuning_tools/
 """
 
+import json
 import logging
 
 from typing import Any
@@ -94,15 +95,17 @@ def get_config_model(self) -> type[BaseModel] | None:
 
     def validate(self, sample: dict) -> bool:
         """Validate that sample has required fields for TRL format."""
-        # Must have messages
-        if "messages" not in sample or not isinstance(sample["messages"], list):
-            return False
+        # Accept either messages format OR agent_cot_tools format
 
-        # Must have at least one message
-        # Return whether there is at least one message
-        # Should have available_tools for tool calling
-        # (though we'll still process samples without it)
-        return len(sample["messages"]) != 0
+        # Check for messages format
+        if "messages" in sample and isinstance(sample["messages"], list):
+            return len(sample["messages"]) > 0
+
+        # Check for agent_cot_tools format (question + tool_used + answer/final_answer)
+        if "question" in sample and "tool_used" in sample:
+            return "answer" in sample or "final_answer" in sample
+
+        return False
 
     def _format_single_sample(self, sample: dict) -> dict | None:
         """
@@ -124,6 +127,10 @@ def _format_single_sample(self, sample: dict) -> dict | None:
             else TRLSFTToolsConfig()
         )
 
+        # Convert agent_cot_tools format to messages format if needed
+        if "messages" not in sample:
+            sample = self._convert_agent_to_messages(sample, config)
+
         # Start with a copy of the sample
         formatted_sample = sample.copy()
 
@@ -172,6 +179,89 @@ def _format_single_sample(self, sample: dict) -> dict | None:
 
         return formatted_sample
 
+    def _convert_agent_to_messages(
+        self, sample: dict, config: TRLSFTToolsConfig
+    ) -> dict:
+        """
+        Convert agent_cot_tools format to messages format.
+
+        Args:
+            sample: Sample in agent_cot_tools format
+            config: Formatter configuration
+
+        Returns:
+            Sample with messages field
+        """
+        messages = []
+
+        # Add system message if configured
+        if config.include_system_prompt:
+            system_content = config.system_prompt_override or (
+                "You are a helpful AI assistant with access to various tools and functions. "
+                "When a user asks a question that requires information or actions you cannot "
+                "directly provide, use the available tools to help answer the question."
+            )
+            messages.append({"role": "system", "content": system_content})
+
+        # Add user question
+        question = sample.get("question", "")
+        messages.append({"role": "user", "content": question})
+
+        # Extract tool usage information
+        tool_used = sample.get("tool_used", "")
+        tool_input = sample.get("tool_input", "{}")
+        tool_output = sample.get("tool_output", "")
+        answer = sample.get("answer") or sample.get("final_answer", "")
+
+        # Parse tool input
+        if isinstance(tool_input, str):
+            try:
+                # First try parsing as standard JSON
+                tool_args = json.loads(tool_input)
+            except json.JSONDecodeError:
+                try:
+                    # Fallback: try replacing single quotes with double quotes
+                    tool_args = json.loads(tool_input.replace("'", '"'))
+                except json.JSONDecodeError:
+                    # Final fallback: wrap in a simple structure
+                    tool_args = {"input": tool_input}
+        else:
+            tool_args = tool_input
+
+        # Add assistant message with tool call
+        # Using OpenAI-compatible function calling format
+        tool_call = {
+            "id": "call_1",  # Placeholder ID
+            "type": "function",
+            "function": {"name": tool_used, "arguments": json.dumps(tool_args)},
+        }
+
+        messages.append(
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [tool_call],
+            }
+        )
+
+        # Add tool response message
+        messages.append(
+            {
+                "role": "tool",
+                "content": str(tool_output),
+                "tool_call_id": "call_1",
+            }
+        )
+
+        # Add final assistant answer
+        messages.append({"role": "assistant", "content": answer})
+
+        # Return sample with messages and preserve available_tools
+        return {
+            "messages": messages,
+            "available_tools": sample.get("available_tools", []),
+        }
+
     def _validate_tool_schemas(self, tools: list[dict]) -> None:
         """
         Validate that tool schemas are properly formatted for TRL.

examples/agent_tool_calling.yaml

@@ -24,12 +24,12 @@ topic_tree:
   topic_prompt: "Real-world scenarios requiring tool usage for task completion"
 
   # Model configuration
-  provider: "gemini"
-  model: "gemini-2.5-flash-lite"
+  provider: "openai"
+  model: "gpt-4o"
   temperature: 0.7
 
   # Tree structure parameters
-  depth: 3
+  depth: 4
   degree: 3
 
   # Save the generated tree
@@ -52,13 +52,10 @@ data_engine:
     Make examples pedagogically valuable for training reasoning skills.
 
   # Model for generating the actual training data
-  provider: "gemini"
-  model: "gemini-2.5-flash-lite"
+  provider: "openai"
+  model: "gpt-4o"
   temperature: 0.8
 
-  # Number of samples to generate
-  num_samples: 50
-
   # Use the agent_cot_hybrid schema for detailed step-by-step reasoning with tools
   conversation_type: "agent_cot_hybrid"
 
@@ -70,7 +67,7 @@ data_engine:
   # tool_registry_path: "examples/custom_tools.yaml"
 
   # Save the generated data
-  save_as: "agent_tool_calling_raw.jsonl"
+  save_as: "agent_tool_calling.jsonl"
 
 # Dataset configuration
 dataset:
@@ -79,8 +76,8 @@ dataset:
 
   creation:
     # Dataset creation parameters
-    num_steps: 2
-    batch_size: 1
+    num_steps: 20
+    batch_size: 4
     sys_msg: true
 
   # Apply the tool-calling formatter to convert to embedded execution format

examples/trl_format_config.yaml

@@ -0,0 +1,10 @@
+dataset:
+dataset:
+  formatters:
+    - name: "trl_sft_tools"
+      template: "builtin://trl_sft_tools"
+      output: "agent_tool_calling_final_trl_sft_tools.jsonl"
+      config:
+        include_system_prompt: true
+        validate_tool_schemas: true
+        remove_available_tools_field: true