Training Pipeline

Author: lukehinds

deepfabric/__init__.py

@@ -17,8 +17,32 @@
 from .generator import DataSetGenerator, DataSetGeneratorConfig
 from .graph import Graph, GraphConfig
 from .hf_hub import HFUploader
+from .pipeline import DeepFabricPipeline
 from .tree import Tree, TreeConfig
 
+# Training module (optional dependencies)
+try:
+    from .training import DeepFabricSFTTrainer, SFTTrainingConfig
+    from .training.sft_config import LoRAConfig, QuantizationConfig
+
+    _has_training = True
+except ImportError:
+    _has_training = False
+    DeepFabricSFTTrainer = None
+    SFTTrainingConfig = None
+    LoRAConfig = None
+    QuantizationConfig = None
+
+# Transformers provider (optional dependencies)
+try:
+    from .llm.transformers_provider import TransformersConfig, TransformersProvider
+
+    _has_transformers = True
+except ImportError:
+    _has_transformers = False
+    TransformersConfig = None
+    TransformersProvider = None
+
 __version__ = "0.1.0"
 
 __all__ = [
@@ -31,7 +55,16 @@
     "Dataset",
     "DeepFabricConfig",
     "HFUploader",
+    "DeepFabricPipeline",
     "cli",
+    # Training (optional)
+    "DeepFabricSFTTrainer",
+    "SFTTrainingConfig",
+    "LoRAConfig",
+    "QuantizationConfig",
+    # Transformers provider (optional)
+    "TransformersConfig",
+    "TransformersProvider",
     # Exceptions
     "DeepFabricError",
     "ConfigurationError",

deepfabric/format_command.py

@@ -52,7 +52,7 @@ def format_command(
         tui.info(f"Loading dataset from Hugging Face repo '{repo}' (split: {hf_split})...")
         try:
             # Bandit nosec, as no digest is set.
-            hf_ds = load_dataset(str(repo), split=hf_split) #  nosec
+            hf_ds = load_dataset(str(repo), split=hf_split)  #  nosec
         except (DatasetNotFoundError, UnexpectedSplitsError) as e:
             msg = (
                 "Failed to load dataset from Hugging Face repo "

deepfabric/llm/client.py

@@ -156,6 +156,41 @@ def make_outlines_model(provider: str, model_name: str, **kwargs) -> Any:
             )
             return outlines.from_openai(client, model_name)
 
+        if provider == "transformers":
+            # Use local HuggingFace Transformers model
+            from .transformers_provider import TransformersProvider  # noqa: PLC0415
+
+            # Create the provider and keep it alive
+            transformers_provider = TransformersProvider(model_name, **kwargs)
+            outlines_model = transformers_provider.get_outlines_model()
+
+            # Create a callable wrapper that uses the Transformers.generate method
+            # This makes the transformers provider work like the API providers
+            def transformers_callable(prompt: str, schema: type[BaseModel], **gen_kwargs):
+                # Import generator module to create logits processor
+                from outlines import generator  # noqa: PLC0415
+
+                # Create JSON schema string from Pydantic model
+                json_schema_str = schema.model_json_schema()
+                # Get JSON string representation
+                import json  # noqa: PLC0415
+
+                json_schema_json = json.dumps(json_schema_str)
+
+                # Convert to logits processor for transformers
+                logits_processor = generator.get_json_schema_logits_processor(
+                    None,  # backend_name (None for auto-detect)
+                    outlines_model,
+                    json_schema_json,
+                )
+                # Generate using the model's generate method
+                return outlines_model.generate(prompt, output_type=logits_processor, **gen_kwargs)
+
+            # Attach the provider to the callable to keep it alive
+            transformers_callable._provider = transformers_provider  # type: ignore[attr-defined]
+
+            return transformers_callable
+
         _raise_unsupported_provider_error(provider)
 
     except DataSetGeneratorError:
@@ -324,7 +359,12 @@ async def _generate_async_with_retry(self, prompt: str, schema: Any, **kwargs) -
         if self.provider == "gemini" and isinstance(schema, type) and issubclass(schema, BaseModel):
             generation_schema = _create_gemini_compatible_schema(schema)
 
-        json_output = await self.async_model(prompt, generation_schema, **kwargs)
+        # Ensure async_model is available; fallback to synchronous generation in a thread if not.
+        model = self.async_model
+        if model is None:
+            return await asyncio.to_thread(self.generate, prompt, schema, **kwargs)
+
+        json_output = await model(prompt, generation_schema, **kwargs)
         # Validate with original schema to ensure proper validation
         return schema.model_validate_json(json_output)
 
@@ -334,6 +374,10 @@ def _convert_generation_params(self, **kwargs) -> dict:
         if self.provider == "gemini" and "max_tokens" in kwargs:
             kwargs["max_output_tokens"] = kwargs.pop("max_tokens")
 
+        # Convert max_tokens to max_new_tokens for Transformers
+        if self.provider == "transformers" and "max_tokens" in kwargs:
+            kwargs["max_new_tokens"] = kwargs.pop("max_tokens")
+
         return kwargs
 
     def __repr__(self) -> str:

deepfabric/llm/rate_limit_config.py

@@ -193,11 +193,59 @@ class OllamaRateLimitConfig(RateLimitConfig):
     )
 
 
+class TransformersRateLimitConfig(RateLimitConfig):
+    """HuggingFace Transformers-specific rate limit configuration.
+
+    Local inference with Transformers doesn't face API rate limits, but
+    may encounter hardware-related failures (CUDA OOM, generation errors).
+    This config uses minimal retries focused on recoverable errors.
+    """
+
+    max_retries: int = Field(
+        default=2,
+        ge=0,
+        le=5,
+        description="Minimal retries for local model inference",
+    )
+    base_delay: float = Field(
+        default=1.0,
+        ge=0.1,
+        le=10.0,
+        description="Base delay for local inference retry",
+    )
+    max_delay: float = Field(
+        default=10.0,
+        ge=1.0,
+        le=60.0,
+        description="Max delay for local inference retry",
+    )
+    backoff_strategy: BackoffStrategy = Field(
+        default=BackoffStrategy.LINEAR,
+        description="Linear backoff for hardware issues",
+    )
+    jitter: bool = Field(
+        default=False,
+        description="No jitter needed for local inference",
+    )
+    respect_retry_after: bool = Field(
+        default=False,
+        description="No retry-after headers from local models",
+    )
+    retry_on_status_codes: set[int] = Field(
+        default_factory=set,
+        description="No HTTP status codes for local inference",
+    )
+    retry_on_exceptions: list[str] = Field(
+        default_factory=lambda: ["cuda", "out of memory", "generation"],
+        description="Exception keywords specific to local model inference",
+    )
+
+
 def get_default_rate_limit_config(provider: str) -> RateLimitConfig:
     """Get the default rate limit configuration for a provider.
 
     Args:
-        provider: Provider name (openai, anthropic, gemini, ollama)
+        provider: Provider name (openai, anthropic, gemini, ollama, transformers)
 
     Returns:
         Provider-specific rate limit configuration with sensible defaults
@@ -207,6 +255,7 @@ def get_default_rate_limit_config(provider: str) -> RateLimitConfig:
         "anthropic": AnthropicRateLimitConfig(),
         "gemini": GeminiRateLimitConfig(),
         "ollama": OllamaRateLimitConfig(),
+        "transformers": TransformersRateLimitConfig(),
     }
     return configs.get(provider, RateLimitConfig())
 
@@ -218,7 +267,7 @@ def create_rate_limit_config(
     """Create a rate limit configuration from a dictionary.
 
     Args:
-        provider: Provider name (openai, anthropic, gemini, ollama)
+        provider: Provider name (openai, anthropic, gemini, ollama, transformers)
         config_dict: Configuration parameters as dictionary
 
     Returns:
@@ -235,6 +284,7 @@ def create_rate_limit_config(
         "anthropic": AnthropicRateLimitConfig,
         "gemini": GeminiRateLimitConfig,
         "ollama": OllamaRateLimitConfig,
+        "transformers": TransformersRateLimitConfig,
     }
 
     config_class = config_classes.get(provider, RateLimitConfig)