Add hydra integration and configuration support in train.py, allowing dynamic model loading and training control. Update requirements.txt to include hydra-core dependency and introduce config.yaml for model parameters and training settings.

conf/config.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - _self_
+
+model_name: "unsloth/SmolLM2-135M-Instruct-bnb-4bit"
+train: false
+output_dir: "final_model"

requirements.txt

@@ -4,6 +4,7 @@ bitsandbytes>=0.45.5
 duckduckgo-search>=8.0.1
 gradio[oauth]>=5.26.0
 hf-xet>=1.0.5
+hydra-core>=1.3.2
 ipywidgets>=8.1.6
 isort>=6.0.1
 jupyter>=1.1.1

train.py

@@ -19,6 +19,9 @@ from datetime import datetime
 from pathlib import Path
 from typing import Union
 
+import hydra
+from omegaconf import DictConfig, OmegaConf
+
 # isort: off
 from unsloth import FastLanguageModel, is_bfloat16_supported  # noqa: E402
 from unsloth.chat_templates import get_chat_template  # noqa: E402
@@ -41,7 +44,6 @@ from transformers import (
 from trl import SFTTrainer
 
 # Configuration
-DEFAULT_MODEL_NAME = "unsloth/SmolLM2-135M-Instruct-bnb-4bit"
 dtype = None  # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
 load_in_4bit = True  # Use 4bit quantization to reduce memory usage
 max_seq_length = 2048  # Auto supports RoPE Scaling internally
@@ -88,7 +90,7 @@ def install_dependencies():
         raise
 
 
-def load_model(model_name: str = DEFAULT_MODEL_NAME) -> tuple[FastLanguageModel, AutoTokenizer]:
+def load_model(model_name: str) -> tuple[FastLanguageModel, AutoTokenizer]:
     """Load and configure the model."""
     logger.info("Loading model and tokenizer...")
     try:
@@ -241,16 +243,18 @@ def create_trainer(
         raise
 
 
-def main():
+@hydra.main(version_base=None, config_path="conf", config_name="config")
+def main(cfg: DictConfig) -> None:
     """Main training function."""
     try:
         logger.info("Starting training process...")
+        logger.info(f"Configuration:\n{OmegaConf.to_yaml(cfg)}")
 
         # Install dependencies
         install_dependencies()
 
         # Load model and tokenizer
-        model, tokenizer = load_model()
+        model, tokenizer = load_model(cfg.model_name)
 
         # Load and prepare dataset
         dataset, tokenizer = load_and_format_dataset(tokenizer)
@@ -258,19 +262,22 @@ def main():
         # Create trainer
         trainer: Trainer = create_trainer(model, tokenizer, dataset)
 
-        # Train
-        logger.info("Starting training...")
-        trainer.train()
-
-        # Save model
-        logger.info("Saving final model...")
-        trainer.save_model("final_model")
-
-        # Print final metrics
-        final_metrics = trainer.state.log_history[-1]
-        logger.info("\nTraining completed!")
-        logger.info(f"Final training loss: {final_metrics.get('loss', 'N/A')}")
-        logger.info(f"Final validation loss: {final_metrics.get('eval_loss', 'N/A')}")
+        # Train if requested
+        if cfg.train:
+            logger.info("Starting training...")
+            trainer.train()
+
+            # Save model
+            logger.info(f"Saving final model to {cfg.output_dir}...")
+            trainer.save_model(cfg.output_dir)
+
+            # Print final metrics
+            final_metrics = trainer.state.log_history[-1]
+            logger.info("\nTraining completed!")
+            logger.info(f"Final training loss: {final_metrics.get('loss', 'N/A')}")
+            logger.info(f"Final validation loss: {final_metrics.get('eval_loss', 'N/A')}")
+        else:
+            logger.info("Training skipped as train=False")
 
     except Exception as e:
         logger.error(f"Error in main training process: {e}")