DeepFabric 4.9.0__py3-none-any.whl → 4.10.0__py3-none-any.whl

deepfabric/config.py

@@ -20,6 +20,7 @@ from .constants import (
 )
 from .exceptions import ConfigurationError
 from .metrics import trace
+from .utils import parse_num_samples
 
 
 def _normalize_reasoning_style(value: str | None) -> str | None:
@@ -131,28 +132,6 @@ class ConversationConfig(BaseModel):
         default=None,
         description="Reasoning style for cot: freetext or agent. Note: 'structured' and 'hybrid' are deprecated.",
     )
-    agent_mode: Literal["single_turn", "multi_turn"] | None = Field(
-        default=None,
-        description="Agent mode: single_turn (one-shot tool use), multi_turn (extended conversations)",
-    )
-    min_turns: int = Field(
-        default=2,
-        ge=1,
-        le=10,
-        description="Minimum conversation turns for multi_turn agent mode",
-    )
-    max_turns: int = Field(
-        default=4,
-        ge=1,
-        le=10,
-        description="Maximum conversation turns for multi_turn agent mode",
-    )
-    min_tool_calls: int = Field(
-        default=2,
-        ge=0,
-        le=20,
-        description="Minimum tool calls before allowing conversation conclusion",
-    )
 
     @field_validator("reasoning_style", mode="before")
     @classmethod
@@ -174,12 +153,6 @@ class ConversationConfig(BaseModel):
                 "Choose from: 'freetext' or 'agent'"
             )
 
-        if self.agent_mode is not None and self.reasoning_style == "freetext":
-            raise ValueError(
-                "reasoning_style='freetext' is not compatible with agent_mode. "
-                "Agent mode requires structured reasoning. Use reasoning_style='agent' instead."
-            )
-
         return self
 
 
@@ -289,22 +262,28 @@ class GenerationConfig(BaseModel):
         default=None, description="Optional LLM configuration overrides for generation"
     )
 
-    @model_validator(mode="after")
-    def validate_agent_requires_tools(self):
-        """Validate that agent_mode requires tools with Spin endpoint."""
-        if self.conversation.agent_mode is not None:
-            if self.tools is None:
-                raise ValueError(
-                    "agent_mode requires tools to be configured. "
-                    "Specify tools.spin_endpoint and optionally tools.available to filter tools."
-                )
-            if not self.tools.spin_endpoint:
-                raise ValueError(
-                    "agent_mode requires a Spin endpoint for tool execution. "
-                    "Set tools.spin_endpoint (e.g., 'http://localhost:3000'). "
-                    "See: cd tools-sdk && spin build && spin up"
-                )
-        return self
+
+class CheckpointConfig(BaseModel):
+    """Configuration for checkpoint-based resume capability.
+
+    Checkpoints allow pausing and resuming long-running dataset generation
+    without losing progress. When enabled, samples are periodically saved
+    to disk and can be resumed if generation is interrupted.
+    """
+
+    interval: int = Field(
+        ...,
+        ge=1,
+        description="Save checkpoint every N samples",
+    )
+    path: str | None = Field(
+        default=None,
+        description="Directory to store checkpoint files. If not specified, uses ~/.deepfabric/checkpoints/{config_hash}/",
+    )
+    retry_failed: bool = Field(
+        default=False,
+        description="When resuming, retry previously failed samples",
+    )
 
 
 class OutputConfig(BaseModel):
@@ -318,10 +297,9 @@ class OutputConfig(BaseModel):
         default=True,
         description="Whether to include system message in output format",
     )
-    num_samples: int = Field(
+    num_samples: int | str = Field(
         default=ENGINE_DEFAULT_NUM_EXAMPLES,
-        ge=1,
-        description="Number of training samples to generate",
+        description="Number of samples: integer, 'auto' (100% of topics), or percentage like '50%'",
     )
     batch_size: int = Field(
         default=ENGINE_DEFAULT_BATCH_SIZE,
@@ -330,6 +308,20 @@ class OutputConfig(BaseModel):
     )
     save_as: str = Field(..., min_length=1, description="Where to save the final dataset")
 
+    # Optional checkpoint configuration (nested inside output)
+    checkpoint: CheckpointConfig | None = Field(
+        None, description="Checkpoint configuration for resumable generation"
+    )
+
+    @field_validator("num_samples", mode="before")
+    @classmethod
+    def validate_num_samples(cls, v: int | str) -> int | str:
+        """Validate num_samples: integer, 'auto', or percentage like '50%'."""
+        result = parse_num_samples(v)
+        if result is None:
+            raise ValueError("num_samples cannot be None")
+        return result
+
 
 class HuggingFaceConfig(BaseModel):
     """Configuration for Hugging Face Hub integration."""
@@ -388,10 +380,6 @@ class EvaluationConfig(BaseModel):
         """Normalize deprecated reasoning_style values."""
         return _normalize_reasoning_style(v)
 
-    agent_mode: Literal["single_turn", "multi_turn"] | None = Field(
-        default=None,
-        description="Agent mode if tools are used",
-    )
     metrics: list[str] = Field(
         default_factory=lambda: [
             "tool_selection_accuracy",
@@ -455,12 +443,6 @@ class EvaluationConfig(BaseModel):
                 "Choose from: 'freetext' or 'agent'"
             )
 
-        if self.agent_mode is not None and self.reasoning_style == "freetext":
-            raise ValueError(
-                "reasoning_style='freetext' is not compatible with agent_mode. "
-                "Agent mode requires structured reasoning. Use reasoning_style='agent' instead."
-            )
-
         return self
 
 
@@ -640,13 +622,17 @@ See documentation for full examples.
             # Conversation config
             "conversation_type": self.generation.conversation.type,
             "reasoning_style": self.generation.conversation.reasoning_style,
-            "agent_mode": self.generation.conversation.agent_mode,
-            "min_turns": self.generation.conversation.min_turns,
-            "max_turns": self.generation.conversation.max_turns,
-            "min_tool_calls": self.generation.conversation.min_tool_calls,
             # Output config
             "sys_msg": self.output.include_system_message,
             "dataset_system_prompt": self.output.system_prompt or self.generation.system_prompt,
+            "output_save_as": self.output.save_as,
+            # Checkpoint config (nested inside output)
+            # Note: checkpoint_path can be None, meaning "auto-resolve" at runtime
+            "checkpoint_interval": self.output.checkpoint.interval if self.output.checkpoint else None,
+            "checkpoint_path": self.output.checkpoint.path if self.output.checkpoint else None,
+            "checkpoint_retry_failed": (
+                self.output.checkpoint.retry_failed if self.output.checkpoint else False
+            ),
         }
 
         # Tool config
@@ -683,6 +669,16 @@ See documentation for full examples.
             "save_as": self.output.save_as,
         }
 
+    def get_checkpoint_config(self) -> dict:
+        """Get checkpoint configuration."""
+        if self.output.checkpoint is None:
+            return {
+                "interval": None,
+                "path": None,  # None means "auto-resolve" at runtime
+                "retry_failed": False,
+            }
+        return self.output.checkpoint.model_dump()
+
     def get_huggingface_config(self) -> dict:
         """Get Hugging Face configuration."""
         return self.huggingface.model_dump() if self.huggingface else {}
@@ -854,10 +850,6 @@ class DataEngineConfig(BaseModel):
     def normalize_reasoning_style(cls, v: str | None) -> str | None:
         return _normalize_reasoning_style(v)
 
-    agent_mode: Literal["single_turn", "multi_turn"] | None = Field(
-        default=None,
-        description="Agent mode for tool use",
-    )
     available_tools: list[str] = Field(
         default_factory=list,
         description="List of tool names available",
@@ -883,14 +875,6 @@ class DataEngineConfig(BaseModel):
                 "Choose from: 'freetext' or 'agent'"
             )
 
-        if self.agent_mode is not None:
-            has_tools = bool(self.available_tools or self.custom_tools)
-            if not has_tools:
-                raise ValueError("agent_mode requires tools to be configured.")
-
-        if self.agent_mode is not None and self.reasoning_style == "freetext":
-            raise ValueError("reasoning_style='freetext' is not compatible with agent_mode.")
-
         return self
 
 

deepfabric/config_manager.py

@@ -65,7 +65,7 @@ def load_config(  # noqa: PLR0913
         mode: Topic generation mode (tree or graph)
         conversation_type: Base conversation type (basic, cot)
         reasoning_style: Reasoning style for cot (freetext, agent)
-        agent_mode: Agent mode (single_turn, multi_turn)
+        agent_mode: [Deprecated] Agent mode (single_turn only, multi_turn no longer supported)
 
     Returns:
         DeepFabricConfig object
@@ -140,7 +140,7 @@ def load_config(  # noqa: PLR0913
             "include_system_message": include_system_message
             if include_system_message is not None
             else True,
-            "num_samples": num_samples or ENGINE_DEFAULT_NUM_EXAMPLES,
+            "num_samples": num_samples if num_samples is not None else ENGINE_DEFAULT_NUM_EXAMPLES,
             "batch_size": batch_size or ENGINE_DEFAULT_BATCH_SIZE,
             "save_as": output_save_as or "dataset.jsonl",
         },
@@ -221,27 +221,29 @@ def apply_cli_overrides(
 
 def get_final_parameters(
     config: DeepFabricConfig,
-    num_samples: int | None = None,
+    num_samples: int | str | None = None,
     batch_size: int | None = None,
     depth: int | None = None,
     degree: int | None = None,
-) -> tuple[int, int, int, int]:
+) -> tuple[int | str, int, int, int]:
     """
     Get final parameters from config and CLI overrides.
 
     Args:
         config: DeepFabricConfig object
-        num_samples: CLI override for num_samples
+        num_samples: CLI override for num_samples (int, "auto", or percentage like "50%")
         batch_size: CLI override for batch_size
         depth: CLI override for depth
         degree: CLI override for degree
 
     Returns:
         Tuple of (num_samples, batch_size, depth, degree)
+        Note: num_samples may be int, "auto", or percentage string
     """
     output_config = config.get_output_config()
 
-    final_num_samples = num_samples or output_config["num_samples"]
+    # Use 'is not None' to allow passing through "auto" or percentage strings
+    final_num_samples = num_samples if num_samples is not None else output_config["num_samples"]
     final_batch_size = batch_size or output_config["batch_size"]
 
     # Get depth and degree from topics config

deepfabric/constants.py

@@ -89,6 +89,12 @@ ERROR_DATASET_FILENAME = "error_dataset.jsonl"
 PARTIAL_TREE_FILENAME = "partial_tree.jsonl"
 FAILED_TREE_SUFFIX = "_failed.jsonl"
 
+# Checkpoint file patterns
+CHECKPOINT_METADATA_SUFFIX = ".checkpoint.json"
+CHECKPOINT_SAMPLES_SUFFIX = ".checkpoint.jsonl"
+CHECKPOINT_FAILURES_SUFFIX = ".checkpoint.failures.jsonl"
+CHECKPOINT_VERSION = 3  # Increment when checkpoint format changes
+
 # Stream simulation defaults
 STREAM_SIM_CHUNK_SIZE = 8  # characters per chunk
 STREAM_SIM_CHUNK_DELAY_MS = 10.0  # milliseconds between chunks

deepfabric/dataset_manager.py

@@ -6,6 +6,7 @@ import traceback
 
 from collections.abc import AsyncIterator
 from datetime import datetime, timezone
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 from datasets import Dataset as HFDataset
@@ -51,6 +52,39 @@ if TYPE_CHECKING:
 DEBUG_MAX_FAILURES_TO_SHOW = 10
 
 
+def resolve_num_samples(num_samples: int | str, topic_count: int) -> int:
+    """Resolve num_samples to an integer based on topic count.
+
+    Args:
+        num_samples: Integer, "auto", or percentage string like "50%"
+        topic_count: Number of available topic paths
+
+    Returns:
+        Resolved integer sample count
+
+    Raises:
+        ConfigurationError: If topic_count is 0 and dynamic sampling is requested
+    """
+    if isinstance(num_samples, int):
+        return num_samples
+
+    if topic_count == 0:
+        raise ConfigurationError(
+            "Cannot use 'auto' or percentage num_samples with empty topic model. "
+            "Ensure topic generation produced paths."
+        )
+
+    if num_samples == "auto":
+        return topic_count
+
+    if isinstance(num_samples, str) and num_samples.endswith("%"):
+        percentage = float(num_samples[:-1]) / 100.0
+        return max(1, int(topic_count * percentage))
+
+    # Fallback - try to parse as int (shouldn't reach here if validated properly)
+    return int(num_samples)
+
+
 async def handle_dataset_events_async(
     generator: AsyncIterator[dict | HFDataset], engine=None, debug: bool = False
 ) -> HFDataset | None:
@@ -80,6 +114,7 @@ async def handle_dataset_events_async(
                         tui.init_status(
                             total_steps=event["num_steps"],
                             total_samples=event["total_samples"],
+                            checkpoint_enabled=event.get("checkpoint_enabled", False),
                         )
 
                         # Build layout with footer card
@@ -112,7 +147,9 @@ async def handle_dataset_events_async(
                         # Footer run status
                         footer_prog = tui.tui.create_footer(layout, title="Run Status")
                         task = footer_prog.add_task(
-                            "Generating dataset samples", total=event["total_samples"]
+                            "Generating dataset samples",
+                            total=event["total_samples"],
+                            completed=event.get("resumed_samples", 0),
                         )
 
                         # Use alternate screen to avoid scroll trails; leave a clean terminal
@@ -130,7 +167,10 @@ async def handle_dataset_events_async(
                         tui.show_generation_header(
                             event["model_name"], event["num_steps"], event["batch_size"]
                         )
-                        simple_task = {"count": 0, "total": event["total_samples"]}
+                        simple_task = {
+                            "count": event.get("resumed_samples", 0),
+                            "total": event["total_samples"],
+                        }
                 elif event["event"] == "step_complete":
                     samples_generated = event.get("samples_generated", 0)
                     if footer_prog and task is not None:
@@ -171,6 +211,40 @@ async def handle_dataset_events_async(
                     total = int(event.get("total_steps", 0))
                     tui.status_step_start(step, total)
 
+                elif event["event"] == "checkpoint_saved":
+                    # Display checkpoint save notification
+                    total_samples = event.get("total_samples", 0)
+                    total_failures = event.get("total_failures", 0)
+                    is_final = event.get("final", False)
+
+                    if footer_prog and task is not None:
+                        # Rich mode: log to events panel and update status
+                        if is_final:
+                            tui.log_event(f"💾 Final checkpoint: {total_samples} samples")
+                        else:
+                            tui.log_event(f"💾 Checkpoint: {total_samples} samples")
+                        tui.status_checkpoint_saved(total_samples)
+                    elif isinstance(simple_task, dict):
+                        # Simple mode: print checkpoint notification
+                        checkpoint_msg = f"Checkpoint saved: {total_samples} samples"
+                        if total_failures > 0:
+                            checkpoint_msg += f" ({total_failures} failures)"
+                        if is_final:
+                            checkpoint_msg = "Final " + checkpoint_msg.lower()
+                        tui.info(checkpoint_msg)
+
+                elif event["event"] == "generation_stopped":
+                    # Graceful stop at checkpoint
+                    if live:
+                        live.stop()
+                    tui.console.print()
+                    tui.success(
+                        f"Gracefully stopped: {event['total_samples']} samples saved to checkpoint"
+                    )
+                    if event.get("total_failures", 0) > 0:
+                        tui.info(f"({event['total_failures']} failures recorded)")
+                    tui.info("Resume with: --resume flag")
+
                 elif event["event"] == "generation_complete":
                     if live:
                         live.stop()
@@ -219,7 +293,7 @@ def create_dataset(
     engine: DataSetGenerator,
     topic_model: "TopicModel",
     config: DeepFabricConfig,
-    num_samples: int | None = None,
+    num_samples: int | str | None = None,
     batch_size: int | None = None,
     include_system_message: bool | None = None,
     provider: str | None = None,  # noqa: ARG001
@@ -234,7 +308,7 @@ def create_dataset(
         engine: DataSetGenerator instance
         topic_model: TopicModel (Tree or Graph) to use for generation
         config: DeepFabricConfig object
-        num_samples: Override for number of samples
+        num_samples: Override for number of samples (int, "auto", or percentage like "50%")
         batch_size: Override for batch size
         include_system_message: Override for including system message
         provider: Override for LLM provider
@@ -268,7 +342,7 @@ async def create_dataset_async(
     engine: DataSetGenerator,
     topic_model: "TopicModel",
     config: DeepFabricConfig,
-    num_samples: int | None = None,
+    num_samples: int | str | None = None,
     batch_size: int | None = None,
     include_system_message: bool | None = None,
     provider: str | None = None,  # noqa: ARG001
@@ -278,15 +352,34 @@ async def create_dataset_async(
 ) -> HFDataset:
     output_config = config.get_output_config()
 
-    final_num_samples = num_samples or output_config["num_samples"]
+    raw_num_samples = num_samples if num_samples is not None else output_config["num_samples"]
     final_batch_size = batch_size or output_config["batch_size"]
 
+    # Resolve "auto" or percentage to actual count based on topic paths
+    topic_count = len(topic_model.get_all_paths())
+    final_num_samples = resolve_num_samples(raw_num_samples, topic_count)
+
+    # Log resolution for dynamic values
+    tui = get_dataset_tui()
+    if isinstance(raw_num_samples, str):
+        tui.info(f"Resolved num_samples: {raw_num_samples} → {final_num_samples} samples")
+
     generation_params = config.get_generation_params(**(generation_overrides or {}))
     final_model = model or generation_params.get("model_name", DEFAULT_MODEL)
 
+    # Convert total samples to number of steps (batches)
+    # The generator expects num_steps where total_samples = num_steps * batch_size
+    import math  # noqa: PLC0415
+
+    final_num_steps = math.ceil(final_num_samples / final_batch_size)
+
+    tui.info(
+        f"Dataset generation: {final_num_samples} samples in {final_num_steps} steps "
+        f"(batch_size={final_batch_size})"
+    )
+
     # Create progress reporter and attach TUI as observer for streaming feedback
     progress_reporter = ProgressReporter()
-    tui = get_dataset_tui()
     progress_reporter.attach(tui)
 
     # Attach progress reporter to engine
@@ -294,7 +387,7 @@ async def create_dataset_async(
 
     try:
         generator = engine.create_data_with_events_async(
-            num_steps=final_num_samples,
+            num_steps=final_num_steps,
             batch_size=final_batch_size,
             topic_model=topic_model,
             model_name=final_model,
@@ -448,6 +541,7 @@ def _strip_nulls(obj: Any) -> Any:
 
 def _save_jsonl_without_nulls(dataset: HFDataset, save_path: str) -> None:
     """Save HF Dataset to JSONL, stripping null values injected by Arrow schema."""
+    Path(save_path).parent.mkdir(parents=True, exist_ok=True)
     with open(save_path, "w") as f:
         for row in dataset:
             cleaned = _strip_nulls(dict(row))
@@ -516,9 +610,11 @@ def save_dataset(
         _save_jsonl_without_nulls(dataset, save_path)
         tui.success(f"Dataset saved to: {save_path}")
 
-        # Save failed samples if engine has any
-        if engine and engine.failed_samples:
-            _save_failed_samples(save_path, engine.failed_samples, tui)
+        # Save failed samples if engine has any (including flushed to checkpoint)
+        if engine:
+            all_failures = engine.get_all_failures()
+            if all_failures:
+                _save_failed_samples(save_path, all_failures, tui)
 
         # Handle automatic uploads if configured
         if config:

deepfabric/evaluation/parser.py

@@ -56,9 +56,9 @@ class GroundTruth(BaseModel):
         default=None,
         description="Reasoning style if cot",
     )
-    agent_mode: Literal["single_turn", "multi_turn"] | None = Field(
+    agent_mode: Literal["single_turn"] | None = Field(
         default=None,
-        description="Agent mode if tools are used",
+        description="Agent mode if tools are used (single_turn only)",
     )
     metadata: dict[str, str | int | float | bool] = Field(
         default_factory=dict,
@@ -77,20 +77,20 @@ class GroundTruthParser:
         self,
         conversation_type: Literal["basic", "cot"],
         reasoning_style: Literal["freetext", "agent", "structured", "hybrid"] | None = None,
-        agent_mode: Literal["single_turn", "multi_turn"] | None = None,
+        agent_mode: Literal["single_turn"] | None = None,
     ):
         """Initialize parser with conversation configuration.
 
         Args:
             conversation_type: Type of conversation (basic, cot)
             reasoning_style: Reasoning style for cot
-            agent_mode: Agent mode if tools are used
+            agent_mode: Agent mode if tools are used (single_turn only)
         """
         self.conversation_type: Literal["basic", "cot"] = conversation_type
         self.reasoning_style: Literal["freetext", "agent", "structured", "hybrid"] | None = (
             reasoning_style
         )
-        self.agent_mode: Literal["single_turn", "multi_turn"] | None = agent_mode
+        self.agent_mode: Literal["single_turn"] | None = agent_mode
 
     def parse(self, conversation: Conversation) -> GroundTruth:
         """Extract ground truth from a conversation sample.
@@ -272,7 +272,7 @@ def parse_batch(
     conversations: list[Conversation],
     conversation_type: Literal["basic", "cot"],
     reasoning_style: Literal["freetext", "agent", "structured", "hybrid"] | None = None,
-    agent_mode: Literal["single_turn", "multi_turn"] | None = None,
+    agent_mode: Literal["single_turn"] | None = None,
 ) -> list[GroundTruth]:
     """Parse a batch of conversations to extract ground truth.
 
@@ -280,7 +280,7 @@ def parse_batch(
         conversations: List of Conversation objects
         conversation_type: Type of conversation
         reasoning_style: Reasoning style if cot
-        agent_mode: Agent mode if tools are used
+        agent_mode: Agent mode if tools are used (single_turn only)
 
     Returns:
         List of GroundTruth objects