Refactor train.py to improve code readability and organization. Adjust logging setup for clarity, streamline dependency installation commands, and enhance dataset splitting and formatting processes. Ensure consistent formatting in log messages and code structure.

train.py

@@ -13,8 +13,8 @@ To run this script:
 2. Run: python train.py
 """
 
-import os
 import logging
+import os
 from datetime import datetime
 from pathlib import Path
 from typing import Union
@@ -39,39 +39,41 @@ dtype = (
 load_in_4bit = True  # Use 4bit quantization to reduce memory usage
 validation_split = 0.1  # 10% of data for validation
 
+
 # Setup logging
 def setup_logging():
     """Configure logging for the training process."""
     # Create logs directory if it doesn't exist
     log_dir = Path("logs")
     log_dir.mkdir(exist_ok=True)
-    
+
     # Create a unique log file name with timestamp
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     log_file = log_dir / f"training_{timestamp}.log"
-    
+
     # Configure logging
     logging.basicConfig(
         level=logging.INFO,
-        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-        handlers=[
-            logging.FileHandler(log_file),
-            logging.StreamHandler()
-        ]
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+        handlers=[logging.FileHandler(log_file), logging.StreamHandler()],
     )
-    
+
     logger = logging.getLogger(__name__)
     logger.info(f"Logging initialized. Log file: {log_file}")
     return logger
 
+
 logger = setup_logging()
 
+
 def install_dependencies():
     """Install required dependencies."""
     logger.info("Installing dependencies...")
     try:
-        os.system('pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"')
-        os.system('pip install --no-deps xformers trl peft accelerate bitsandbytes')
+        os.system(
+            'pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"'
+        )
+        os.system("pip install --no-deps xformers trl peft accelerate bitsandbytes")
         logger.info("Dependencies installed successfully")
     except Exception as e:
         logger.error(f"Error installing dependencies: {e}")
@@ -133,7 +135,9 @@ def load_and_format_dataset(
 
         # Split into train and validation sets
         dataset = dataset.train_test_split(test_size=validation_split, seed=3407)
-        logger.info(f"Dataset split into train ({len(dataset['train'])} examples) and validation ({len(dataset['test'])} examples) sets")
+        logger.info(
+            f"Dataset split into train ({len(dataset['train'])} examples) and validation ({len(dataset['test'])} examples) sets"
+        )
 
         # Configure chat template
         tokenizer = get_chat_template(
@@ -160,10 +164,14 @@ def load_and_format_dataset(
             return {"text": texts}
 
         # Apply formatting to both train and validation sets
-        dataset = DatasetDict({
-            "train": dataset["train"].map(formatting_prompts_func, batched=True),
-            "validation": dataset["test"].map(formatting_prompts_func, batched=True)
-        })
+        dataset = DatasetDict(
+            {
+                "train": dataset["train"].map(formatting_prompts_func, batched=True),
+                "validation": dataset["test"].map(
+                    formatting_prompts_func, batched=True
+                ),
+            }
+        )
         logger.info("Dataset formatting completed successfully")
 
         return dataset, tokenizer
@@ -226,33 +234,33 @@ def main():
     """Main training function."""
     try:
         logger.info("Starting training process...")
-        
+
         # Install dependencies
         install_dependencies()
-        
+
         # Load model and tokenizer
         model, tokenizer = load_model()
-        
+
         # Load and prepare dataset
         dataset, tokenizer = load_and_format_dataset(tokenizer)
-        
+
         # 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')}")
-        
+
     except Exception as e:
         logger.error(f"Error in main training process: {e}")
         raise