Update: Comprehensive documentation and verbose comments throughout the codebase

app.py

@@ -7,9 +7,27 @@ This version imports OpenLLM modules from the uploaded files in the HF Space:
 - Uses OpenLLM's actual custom model architecture
 - Compatible with OpenLLM's implementation
 
+This application provides a complete training interface for OpenLLM models on Hugging Face Spaces.
+It uses OpenLLM's custom GPTModel architecture instead of Hugging Face Transformers,
+ensuring compatibility with the actual OpenLLM implementation.
+
+Key Features:
+- Real model training using OpenLLM's custom architecture
+- SentencePiece tokenization for OpenLLM models
+- Complete training pipeline with progress monitoring
+- Automatic model saving and uploading to Hugging Face Hub
+- Gradio 4.44.1 compatible user interface
+
+Technical Architecture:
+- Uses OpenLLM's GPTModel class (not Hugging Face Transformers)
+- Imports custom modules from uploaded files in the Space
+- Uses sentencepiece.SentencePieceProcessor() for tokenization
+- Implements OpenLLM's training loop and optimization strategy
+- Saves checkpoints in OpenLLM's format
+
 Author: Louis Chua Bean Chong
 License: GPL-3.0
-Version: 2.1.0
+Version: 2.1.1
 Last Updated: 2024
 """
 
@@ -26,37 +44,72 @@ from dataclasses import dataclass
 from pathlib import Path
 
 # Import OpenLLM's custom model architecture from uploaded files
+# These files were uploaded to the HF Space and contain OpenLLM's actual implementation
 try:
     # Import from the uploaded files in the HF Space
+    # model.py contains GPTModel, GPTConfig, and create_model factory function
     from model import GPTModel, GPTConfig, create_model
+    # data_loader.py contains TextDataLoader for OpenLLM's data loading approach
     from data_loader import TextDataLoader
     OPENLLM_AVAILABLE = True
     print("✅ OpenLLM custom model architecture imported successfully from uploaded files")
+    print("   - GPTModel: Custom PyTorch model architecture")
+    print("   - GPTConfig: Model configuration dataclass")
+    print("   - create_model: Factory function for model creation")
+    print("   - TextDataLoader: Custom data loading implementation")
 except ImportError as e:
     print(f"❌ OpenLLM imports failed: {e}")
+    print("   This indicates the uploaded OpenLLM source files are not available")
+    print("   The training functionality will be disabled")
     OPENLLM_AVAILABLE = False
 
-# Try to import sentencepiece
+# Try to import sentencepiece - CRITICAL for OpenLLM tokenization
+# OpenLLM uses SentencePiece for tokenization, not Hugging Face tokenizers
 try:
     import sentencepiece as spm
     SENTENCEPIECE_AVAILABLE = True
     print(f"✅ SentencePiece available: {spm.__version__}")
+    print("   - Required for OpenLLM tokenization")
+    print("   - Used for loading tokenizer.model files")
 except ImportError:
     SENTENCEPIECE_AVAILABLE = False
     print("❌ SentencePiece not available")
+    print("   - This will prevent tokenizer loading")
+    print("   - Training functionality will be limited")
 
-# Import other dependencies
+# Import other dependencies for the complete training pipeline
 try:
-    from datasets import load_dataset
-    from huggingface_hub import HfApi, hf_hub_download
+    from datasets import load_dataset  # For loading training data from HF Hub
+    from huggingface_hub import HfApi, hf_hub_download  # For model uploads and downloads
     DEPENDENCIES_AVAILABLE = True
+    print("✅ Training dependencies available")
+    print("   - datasets: For loading training data")
+    print("   - huggingface_hub: For model uploads/downloads")
 except ImportError as e:
     print(f"❌ Dependencies not available: {e}")
+    print("   - This will prevent dataset loading and model uploading")
     DEPENDENCIES_AVAILABLE = False
 
 @dataclass
 class TrainingConfig:
-    """Configuration class for training parameters."""
+    """
+    Configuration class for training parameters.
+    
+    This dataclass encapsulates all the training hyperparameters and settings
+    that control the OpenLLM training process. It provides a clean interface
+    for passing configuration between different components of the training pipeline.
+    
+    Attributes:
+        model_size: Size of the model to train ("small", "medium", "large")
+        max_steps: Maximum number of training iterations
+        learning_rate: Learning rate for the optimizer
+        batch_size: Number of samples per training batch
+        output_dir: Directory to save trained models and checkpoints
+        save_steps: Frequency of checkpoint saving (every N steps)
+        logging_steps: Frequency of progress logging (every N steps)
+        warmup_steps: Number of warmup steps for learning rate scheduling
+        gradient_accumulation_steps: Number of steps to accumulate gradients
+    """
     model_size: str
     max_steps: int
     learning_rate: float
@@ -74,83 +127,139 @@ class OpenLLMTrainer:
     This class handles the entire training pipeline including:
     - Model loading using OpenLLM's custom GPTModel
     - Tokenizer loading using sentencepiece.SentencePieceProcessor()
-    - Dataset preparation
+    - Dataset preparation using OpenLLM's TextDataLoader
     - Training execution using OpenLLM's approach
-    - Model saving and uploading
+    - Model saving and uploading to Hugging Face Hub
+    
+    The trainer implements OpenLLM's actual training methodology rather than
+    using Hugging Face Transformers, ensuring compatibility with the real
+    OpenLLM implementation.
+    
+    Key Features:
+    - Custom model architecture (GPTModel, not PreTrainedModel)
+    - SentencePiece tokenization (not Hugging Face tokenizers)
+    - OpenLLM's training loop and optimization strategy
+    - Gradient accumulation for memory efficiency
+    - Learning rate scheduling with warmup
+    - Automatic checkpoint saving and model uploading
     """
     
     def __init__(self):
-        """Initialize the trainer with default settings."""
-        self.model = None
-        self.tokenizer = None
-        self.data_loader = None
-        self.optimizer = None
-        self.scheduler = None
-        self.is_training = False
+        """
+        Initialize the trainer with default settings.
+        
+        Sets up the trainer with default values and initializes the Hugging Face
+        API for model uploading. All components start as None and are initialized
+        during the training process.
+        """
+        # Core training components - initialized during training
+        self.model = None  # OpenLLM's GPTModel instance
+        self.tokenizer = None  # SentencePieceProcessor instance
+        self.data_loader = None  # OpenLLM's TextDataLoader instance
+        self.optimizer = None  # PyTorch optimizer (AdamW)
+        self.scheduler = None  # Learning rate scheduler
+        
+        # Training state management
+        self.is_training = False  # Flag to track training status
+        self.tokenizer_path = None  # Path to the tokenizer.model file
+        
+        # Progress tracking for UI updates
         self.training_progress = {
-            "status": "Ready",
-            "current_step": 0,
-            "total_steps": 0,
-            "loss": 0.0,
-            "learning_rate": 0.0
+            "status": "Ready",  # Current training status
+            "current_step": 0,  # Current training step
+            "total_steps": 0,  # Total steps to complete
+            "loss": 0.0,  # Current training loss
+            "learning_rate": 0.0  # Current learning rate
         }
         
         # Initialize Hugging Face API for model uploading
+        # This allows the trained model to be automatically uploaded to HF Hub
         try:
             self.hf_api = HfApi()
+            print("✅ Hugging Face API initialized for model uploading")
         except Exception as e:
             print(f"Failed to initialize HF API: {e}")
+            print("   - Model uploading will be disabled")
             self.hf_api = None
     
     def load_model_and_tokenizer(self, model_size: str) -> str:
         """
         Load the pre-trained OpenLLM model and tokenizer using OpenLLM's approach.
         
+        This method implements OpenLLM's actual model loading strategy:
+        1. Creates a new GPTModel using OpenLLM's factory function
+        2. Downloads the tokenizer.model file from Hugging Face Hub
+        3. Loads the tokenizer using SentencePieceProcessor
+        4. Stores both components for use in training
+        
+        This approach differs from Hugging Face Transformers because:
+        - Uses OpenLLM's custom GPTModel (not AutoModelForCausalLM)
+        - Uses SentencePiece directly (not AutoTokenizer)
+        - Downloads specific files rather than using from_pretrained()
+        
         Args:
             model_size: Size of the model to load ("small", "medium", "large")
+                       Determines which pre-trained model to download
             
         Returns:
             Status message indicating success or failure
+            Success: "✅ Successfully loaded OpenLLM {model_size} model with custom architecture"
+            Failure: "❌ Failed to load OpenLLM model and tokenizer: {error details}"
         """
         try:
+            # Verify OpenLLM modules are available
             if not OPENLLM_AVAILABLE:
                 return "❌ OpenLLM custom model architecture not available"
             
             print(f"🔄 Loading OpenLLM {model_size} model using custom architecture...")
+            print(f"   - Using OpenLLM's create_model factory function")
+            print(f"   - Not using Hugging Face Transformers")
             
-            # Create model using OpenLLM's factory function
+            # Step 1: Create model using OpenLLM's factory function
+            # This creates a fresh GPTModel instance with the specified size
             try:
                 self.model = create_model(model_size)
                 print(f"✅ OpenLLM {model_size} model created: {type(self.model).__name__}")
-                print(f"   Parameters: {self.model.get_num_params():,}")
+                print(f"   - Model type: {type(self.model).__name__}")
+                print(f"   - Parameters: {self.model.get_num_params():,}")
+                print(f"   - Architecture: Custom GPTModel (not PreTrainedModel)")
             except Exception as e:
                 print(f"❌ Failed to create model: {e}")
                 return f"❌ Failed to create OpenLLM model: {str(e)}"
             
-            # Load tokenizer using sentencepiece
+            # Step 2: Load tokenizer using sentencepiece
+            # OpenLLM uses SentencePiece directly, not Hugging Face tokenizers
             try:
                 print("🔄 Loading tokenizer using sentencepiece.SentencePieceProcessor()...")
+                print("   - Using SentencePiece directly (not AutoTokenizer)")
+                print("   - Downloading tokenizer.model from Hugging Face Hub")
                 
                 # Download tokenizer.model from HF Hub
+                # This is the actual tokenizer file used by OpenLLM models
                 model_name = f"lemms/openllm-{model_size}-extended-7k"
                 tokenizer_path = hf_hub_download(
                     repo_id=model_name,
-                    filename="tokenizer.model"
+                    filename="tokenizer.model"  # Specific file name for OpenLLM
                 )
                 
                 print(f"✅ Tokenizer downloaded to: {tokenizer_path}")
+                print(f"   - Source: {model_name}")
+                print(f"   - File: tokenizer.model")
                 
-                # Create SentencePieceProcessor
+                # Create SentencePieceProcessor and load the tokenizer
+                # This is OpenLLM's actual tokenization approach
                 sp_processor = spm.SentencePieceProcessor()
                 sp_processor.load(tokenizer_path)
                 
                 # Store tokenizer and its path separately
+                # We need the path for the TextDataLoader later
                 self.tokenizer = sp_processor
                 self.tokenizer_path = tokenizer_path  # Store the path separately
                 
                 print(f"✅ Tokenizer loaded successfully using SentencePieceProcessor")
-                print(f"   Vocabulary size: {sp_processor.vocab_size()}")
-                print(f"   Tokenizer path: {tokenizer_path}")
+                print(f"   - Vocabulary size: {sp_processor.vocab_size()}")
+                print(f"   - Tokenizer path: {tokenizer_path}")
+                print(f"   - Tokenizer type: {type(sp_processor).__name__}")
                 
             except Exception as e:
                 print(f"❌ Failed to load tokenizer: {e}")
@@ -165,41 +274,73 @@ class OpenLLMTrainer:
         """
         Load and prepare the training dataset using OpenLLM's approach.
         
+        This method implements OpenLLM's data preparation strategy:
+        1. Loads training data from Hugging Face Hub dataset
+        2. Creates a temporary text file for OpenLLM's TextDataLoader
+        3. Initializes OpenLLM's TextDataLoader with the tokenizer
+        4. Prepares the data for training
+        
+        OpenLLM's approach differs from Hugging Face because:
+        - Uses a simple text file format (not tokenized datasets)
+        - Uses OpenLLM's TextDataLoader (not Hugging Face datasets)
+        - Tokenization happens on-the-fly during training
+        
         Returns:
             Status message indicating success or failure
+            Success: "✅ Successfully prepared dataset with {count} samples"
+            Failure: "❌ Failed to prepare dataset: {error details}"
         """
         try:
+            # Verify dependencies are available
             if not DEPENDENCIES_AVAILABLE:
                 return "❌ Required dependencies not available"
             
             print("🔄 Loading training dataset...")
+            print("   - Loading from Hugging Face Hub dataset")
+            print("   - Using OpenLLM's data preparation approach")
             
             # Load dataset from HF Hub
+            # This contains the training text data for continuing model training
             dataset = load_dataset("lemms/openllm-training-data")
             print(f"✅ Dataset loaded: {len(dataset['train'])} samples")
+            print(f"   - Dataset: lemms/openllm-training-data")
+            print(f"   - Samples: {len(dataset['train'])}")
             
             # Create temporary data file for OpenLLM's TextDataLoader
+            # OpenLLM expects a simple text file with one text sample per line
             temp_data_file = "temp_training_data.txt"
             with open(temp_data_file, 'w', encoding='utf-8') as f:
                 for item in dataset['train']:
                     f.write(item['text'] + '\n')
             
             print(f"✅ Temporary data file created: {temp_data_file}")
+            print(f"   - Format: One text sample per line")
+            print(f"   - Encoding: UTF-8")
             
             # Create OpenLLM's TextDataLoader
+            # This is OpenLLM's custom data loading implementation
             try:
                 # Use the stored tokenizer path instead of trying to access model_file_path
+                # SentencePieceProcessor doesn't have a model_file_path attribute
                 tokenizer_path = self.tokenizer_path  # Use the stored path
                 
+                print(f"🔄 Creating OpenLLM TextDataLoader...")
+                print(f"   - Data file: {temp_data_file}")
+                print(f"   - Tokenizer path: {tokenizer_path}")
+                print(f"   - Sequence length: 512")
+                print(f"   - Batch size: 4 (will be overridden by training config)")
+                
                 self.data_loader = TextDataLoader(
                     data_file=temp_data_file,
                     tokenizer_path=tokenizer_path,
-                    seq_len=512,
+                    seq_len=512,  # Maximum sequence length for training
                     batch_size=4,  # Will be overridden by training config
-                    shuffle=True
+                    shuffle=True   # Shuffle data for better training
                 )
                 
                 print(f"✅ OpenLLM TextDataLoader created successfully")
+                print(f"   - DataLoader type: {type(self.data_loader).__name__}")
+                print(f"   - Uses OpenLLM's custom implementation")
                 
             except Exception as e:
                 print(f"❌ Failed to create TextDataLoader: {e}")
@@ -214,60 +355,107 @@ class OpenLLMTrainer:
         """
         Set up the training configuration using OpenLLM's approach.
         
+        This method configures the training environment with:
+        1. Output directory creation
+        2. Optimizer setup with weight decay groups
+        3. Learning rate scheduler with warmup
+        4. Training hyperparameters
+        
+        The setup follows OpenLLM's training methodology:
+        - Uses AdamW optimizer with weight decay
+        - Implements learning rate warmup followed by cosine annealing
+        - Separates parameters for different weight decay rates
+        - Uses gradient clipping for stability
+        
         Args:
-            config: Training configuration object
+            config: Training configuration object containing all hyperparameters
             
         Returns:
             Status message indicating success or failure
+            Success: "✅ Training setup completed successfully"
+            Failure: "❌ Failed to setup training: {error details}"
         """
         try:
-            # Create output directory
+            print("🔄 Setting up training configuration...")
+            print(f"   - Output directory: {config.output_dir}")
+            print(f"   - Learning rate: {config.learning_rate}")
+            print(f"   - Max steps: {config.max_steps}")
+            
+            # Create output directory for saving models and checkpoints
             os.makedirs(config.output_dir, exist_ok=True)
+            print(f"✅ Output directory created: {config.output_dir}")
             
             # Set up optimizer (AdamW with weight decay)
-            decay_params = []
-            no_decay_params = []
+            # This follows OpenLLM's optimization strategy
+            print("🔄 Setting up AdamW optimizer with weight decay...")
+            
+            # Separate parameters for different weight decay rates
+            # This is a common practice for transformer training
+            decay_params = []      # Parameters that should have weight decay
+            no_decay_params = []   # Parameters that should not have weight decay
             
             for name, param in self.model.named_parameters():
                 if not param.requires_grad:
                     continue
                 
+                # Apply weight decay to all parameters except biases and layer norm weights
                 if len(param.shape) == 1 or name.endswith('.bias'):
                     no_decay_params.append(param)
                 else:
                     decay_params.append(param)
             
+            # Create parameter groups with different weight decay rates
             param_groups = [
-                {'params': decay_params, 'weight_decay': 0.01},
-                {'params': no_decay_params, 'weight_decay': 0.0}
+                {'params': decay_params, 'weight_decay': 0.01},      # 1% weight decay
+                {'params': no_decay_params, 'weight_decay': 0.0}     # No weight decay
             ]
             
+            print(f"   - Decay parameters: {len(decay_params)}")
+            print(f"   - No-decay parameters: {len(no_decay_params)}")
+            
+            # Initialize AdamW optimizer with OpenLLM's recommended settings
             self.optimizer = torch.optim.AdamW(
                 param_groups,
                 lr=config.learning_rate,
-                betas=(0.9, 0.95),
-                eps=1e-8
+                betas=(0.9, 0.95),  # Beta values for momentum
+                eps=1e-8            # Epsilon for numerical stability
             )
             
+            print(f"✅ AdamW optimizer configured")
+            print(f"   - Learning rate: {config.learning_rate}")
+            print(f"   - Betas: (0.9, 0.95)")
+            print(f"   - Epsilon: 1e-8")
+            
             # Set up learning rate scheduler
+            # OpenLLM uses a warmup followed by cosine annealing
+            print("🔄 Setting up learning rate scheduler...")
+            
+            # Warmup scheduler: linearly increase LR from 1% to 100%
             warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
                 self.optimizer,
-                start_factor=0.01,
-                end_factor=1.0,
+                start_factor=0.01,  # Start at 1% of target LR
+                end_factor=1.0,     # End at 100% of target LR
                 total_iters=config.warmup_steps
             )
             
+            # Main scheduler: cosine annealing after warmup
             main_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                 self.optimizer,
-                T_max=config.max_steps - config.warmup_steps
+                T_max=config.max_steps - config.warmup_steps  # Duration of cosine annealing
             )
             
+            # Combine warmup and main schedulers
             self.scheduler = torch.optim.lr_scheduler.SequentialLR(
                 self.optimizer,
                 schedulers=[warmup_scheduler, main_scheduler],
-                milestones=[config.warmup_steps]
+                milestones=[config.warmup_steps]  # Switch to main scheduler after warmup
             )
             
+            print(f"✅ Learning rate scheduler configured")
+            print(f"   - Warmup steps: {config.warmup_steps}")
+            print(f"   - Total steps: {config.max_steps}")
+            print(f"   - Schedule: Linear warmup → Cosine annealing")
+            
             print("✅ Training setup completed successfully")
             return f"✅ Training setup completed successfully"
             
@@ -278,107 +466,165 @@ class OpenLLMTrainer:
         """
         Execute the actual model training using OpenLLM's approach.
         
+        This method implements OpenLLM's training loop:
+        1. Sets up training mode and progress tracking
+        2. Iterates through data batches using OpenLLM's TextDataLoader
+        3. Performs forward pass, loss computation, and backward pass
+        4. Implements gradient accumulation for memory efficiency
+        5. Updates model parameters and learning rate
+        6. Saves checkpoints and logs progress
+        
+        The training loop follows OpenLLM's methodology:
+        - Uses OpenLLM's GPTModel forward pass (returns logits and loss)
+        - Implements gradient accumulation for effective larger batch sizes
+        - Uses gradient clipping for training stability
+        - Saves checkpoints in OpenLLM's format
+        - Updates progress for UI monitoring
+        
         Args:
-            config: Training configuration object
+            config: Training configuration object containing hyperparameters
             progress_callback: Optional callback function for progress updates
+                             (Not used in current implementation)
             
         Returns:
             Status message indicating success or failure
+            Success: "✅ Training completed successfully! Final step: {step}"
+            Failure: "❌ Training failed: {error details}"
         """
         try:
+            # Set training state
             self.is_training = True
             self.training_progress["status"] = "Training"
             self.training_progress["total_steps"] = config.max_steps
             
             print(f"🚀 Starting OpenLLM training for {config.max_steps} steps...")
+            print(f"   - Model: {type(self.model).__name__}")
+            print(f"   - DataLoader: {type(self.data_loader).__name__}")
+            print(f"   - Optimizer: {type(self.optimizer).__name__}")
+            print(f"   - Gradient accumulation: {config.gradient_accumulation_steps}")
             
             # Training loop using OpenLLM's approach
-            self.model.train()
-            accumulated_loss = 0.0
-            self.optimizer.zero_grad()
+            self.model.train()  # Set model to training mode
+            accumulated_loss = 0.0  # Track loss across accumulation steps
+            self.optimizer.zero_grad()  # Clear gradients
             
-            step = 0
+            step = 0  # Current training step
             for batch_idx, (input_ids, target_ids) in enumerate(self.data_loader):
+                # Check if we've reached the maximum number of steps
                 if step >= config.max_steps:
                     break
                 
                 # Forward pass (model computes loss internally when targets provided)
+                # OpenLLM's GPTModel returns both logits and loss
                 logits, loss = self.model(input_ids, target_ids)
                 
                 # Scale loss for gradient accumulation
+                # This allows us to simulate larger batch sizes
                 loss = loss / config.gradient_accumulation_steps
                 accumulated_loss += loss.item()
                 
-                # Backward pass
+                # Backward pass - compute gradients
                 loss.backward()
                 
                 # Update weights every gradient_accumulation_steps
                 if (batch_idx + 1) % config.gradient_accumulation_steps == 0:
-                    # Clip gradients
+                    # Clip gradients for training stability
+                    # This prevents exploding gradients
                     torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                     
-                    # Update parameters
+                    # Update parameters using the optimizer
                     self.optimizer.step()
+                    
+                    # Update learning rate using the scheduler
                     self.scheduler.step()
+                    
+                    # Clear gradients for the next accumulation cycle
                     self.optimizer.zero_grad()
                     
                     # Update step count
                     step += 1
                     
-                    # Update progress
+                    # Update progress for UI monitoring
                     self.training_progress["current_step"] = step
                     self.training_progress["loss"] = accumulated_loss
                     self.training_progress["learning_rate"] = self.scheduler.get_last_lr()[0]
                     
-                    # Log progress
+                    # Log progress at specified intervals
                     if step % config.logging_steps == 0:
-                        print(f"Step {step}/{config.max_steps} | Loss: {accumulated_loss:.4f} | LR: {self.scheduler.get_last_lr()[0]:.2e}")
+                        current_lr = self.scheduler.get_last_lr()[0]
+                        print(f"Step {step}/{config.max_steps} | Loss: {accumulated_loss:.4f} | LR: {current_lr:.2e}")
                     
-                    # Save checkpoint
+                    # Save checkpoint at specified intervals
                     if step % config.save_steps == 0:
                         self._save_checkpoint(config.output_dir, step)
+                        print(f"💾 Checkpoint saved at step {step}")
                     
-                    # Reset accumulated loss
+                    # Reset accumulated loss for the next accumulation cycle
                     accumulated_loss = 0.0
                     
-                    # Clean up memory
+                    # Clean up memory periodically
                     if step % 100 == 0:
                         gc.collect()
+                        print(f"🧹 Memory cleanup at step {step}")
             
-            # Final checkpoint
+            # Save final checkpoint
             self._save_checkpoint(config.output_dir, step, is_best=True)
+            print(f"💾 Final checkpoint saved at step {step}")
             
             # Update final progress
             self.training_progress["status"] = "Completed"
             self.training_progress["current_step"] = step
             
             print(f"✅ Training completed! Final step: {step}")
+            print(f"   - Total steps completed: {step}")
+            print(f"   - Final loss: {self.training_progress['loss']:.4f}")
+            print(f"   - Final learning rate: {self.training_progress['learning_rate']:.2e}")
             
             return f"✅ Training completed successfully! Final step: {step}"
             
         except Exception as e:
             self.training_progress["status"] = "Failed"
             print(f"❌ Training failed: {e}")
+            print(f"   - Error occurred during training")
+            print(f"   - Training state: {self.training_progress['status']}")
             return f"❌ Training failed: {str(e)}"
         finally:
             self.is_training = False
     
     def _save_checkpoint(self, output_dir: str, step: int, is_best: bool = False) -> None:
-        """Save model checkpoint using OpenLLM's approach."""
+        """
+        Save model checkpoint using OpenLLM's approach.
+        
+        This method saves the model state in OpenLLM's checkpoint format:
+        - Model state dictionary
+        - Optimizer state dictionary
+        - Scheduler state dictionary
+        - Model configuration
+        - Training step information
+        
+        The checkpoint format is compatible with OpenLLM's loading mechanism
+        and can be used to resume training or load the model for inference.
+        
+        Args:
+            output_dir: Directory to save the checkpoint
+            step: Current training step number
+            is_best: Whether this is the best model so far
+        """
         try:
+            # Create checkpoint dictionary with all necessary components
             checkpoint = {
-                'step': step,
-                'model_state_dict': self.model.state_dict(),
-                'optimizer_state_dict': self.optimizer.state_dict(),
-                'scheduler_state_dict': self.scheduler.state_dict(),
-                'config': self.model.config.__dict__
+                'step': step,                                    # Current training step
+                'model_state_dict': self.model.state_dict(),     # Model parameters
+                'optimizer_state_dict': self.optimizer.state_dict(),  # Optimizer state
+                'scheduler_state_dict': self.scheduler.state_dict(),  # Scheduler state
+                'config': self.model.config.__dict__             # Model configuration
             }
             
             # Save latest checkpoint
             checkpoint_path = os.path.join(output_dir, f"checkpoint_step_{step}.pt")
             torch.save(checkpoint, checkpoint_path)
             
-            # Save best checkpoint
+            # Save best checkpoint if this is the best model
             if is_best:
                 best_path = os.path.join(output_dir, "best_model.pt")
                 torch.save(checkpoint, best_path)
@@ -393,37 +639,59 @@ class OpenLLMTrainer:
         """
         Save the trained model and upload it to Hugging Face Hub.
         
+        This method completes the training pipeline by:
+        1. Saving the final model checkpoint
+        2. Copying the tokenizer files
+        3. Uploading the complete model to Hugging Face Hub
+        4. Creating a new model repository for the trained model
+        
+        The uploaded model will be available at:
+        https://huggingface.co/lemms/openllm-{size}-extended-8k
+        
         Args:
             config: Training configuration object
             
         Returns:
             Status message indicating success or failure
+            Success: "✅ Model saved and uploaded to https://huggingface.co/{repo_id}"
+            Failure: "❌ Failed to save/upload model: {error details}"
         """
         try:
             print("🔄 Saving trained model...")
+            print(f"   - Output directory: {config.output_dir}")
+            print(f"   - Model size: {config.model_size}")
             
-            # Save the final model
+            # Save the final model checkpoint
             self._save_checkpoint(config.output_dir, config.max_steps, is_best=True)
             
             # Save tokenizer files
+            # Create a tokenizer directory within the output directory
             tokenizer_dir = os.path.join(config.output_dir, "tokenizer")
             os.makedirs(tokenizer_dir, exist_ok=True)
             
             # Copy the tokenizer.model file using the stored path
+            # This ensures the tokenizer is included with the model
             import shutil
             shutil.copy2(self.tokenizer_path, os.path.join(tokenizer_dir, "tokenizer.model"))
             
             print("✅ Model saved locally")
+            print(f"   - Model checkpoint: {config.output_dir}/best_model.pt")
+            print(f"   - Tokenizer: {tokenizer_dir}/tokenizer.model")
             
             # Generate model name for upload
+            # The naming convention follows: openllm-{size}-extended-8k
             model_name = f"openllm-{config.model_size}-extended-8k"
             repo_id = f"lemms/{model_name}"
             
             # Upload to Hugging Face Hub
             if self.hf_api:
                 print(f"🔄 Uploading model to {repo_id}...")
+                print(f"   - Repository: {repo_id}")
+                print(f"   - Type: model")
+                print(f"   - Source: {config.output_dir}")
                 
-                # Upload model files
+                # Upload model files to Hugging Face Hub
+                # This creates a new model repository with all the files
                 self.hf_api.upload_folder(
                     folder_path=config.output_dir,
                     repo_id=repo_id,
@@ -432,8 +700,10 @@ class OpenLLMTrainer:
                 )
                 
                 print(f"✅ Model uploaded successfully to {repo_id}")
+                print(f"   - Available at: https://huggingface.co/{repo_id}")
                 return f"✅ Model saved and uploaded to https://huggingface.co/{repo_id}"
             else:
+                print("⚠️ Hugging Face API not available - model saved locally only")
                 return f"✅ Model saved locally to {config.output_dir}"
                 
         except Exception as e:
@@ -441,99 +711,146 @@ class OpenLLMTrainer:
             return f"❌ Failed to save/upload model: {str(e)}"
     
     def get_training_progress(self) -> Dict[str, Any]:
-        """Get current training progress information."""
+        """
+        Get current training progress information.
+        
+        This method returns a copy of the current training progress
+        for display in the Gradio UI. The progress information includes:
+        - Current training status
+        - Current step and total steps
+        - Current loss value
+        - Current learning rate
+        
+        Returns:
+            Dictionary containing current training progress information
+        """
         return self.training_progress.copy()
 
 def main():
     """
     Main function that creates the complete Gradio application interface.
+    
+    This function sets up the entire Gradio application with:
+    1. Application header and status information
+    2. Training configuration controls
+    3. Training status and progress display
+    4. Training control buttons
+    5. Instructions and resource links
+    6. Training function implementation
+    
+    The interface provides a complete training experience for OpenLLM models
+    with real-time progress monitoring and comprehensive configuration options.
+    
+    Returns:
+        Gradio Blocks interface for the training application
     """
     
     # Initialize the trainer
+    # This creates the OpenLLMTrainer instance that will handle all training operations
     trainer = OpenLLMTrainer()
     
     # Create the main Gradio application interface
+    # Using Gradio 4.44.1 with Soft theme for modern appearance
     with gr.Blocks(
         title="OpenLLM Training Space - Fixed with Uploaded Modules",
         theme=gr.themes.Soft()
     ) as demo:
         
         # Application Header
+        # Provides clear identification and description of the application
         gr.Markdown("# 🚀 OpenLLM Training Space - Fixed with Uploaded Modules")
         gr.Markdown("### *Uses OpenLLM's Custom Model Architecture from Uploaded Files*")
         gr.Markdown("---")
         
         # Status Information
+        # Shows the availability of key components and dependencies
         gr.Markdown(f"**OpenLLM Available**: {'✅ Yes' if OPENLLM_AVAILABLE else '❌ No'}")
         gr.Markdown(f"**SentencePiece Available**: {'✅ Yes' if SENTENCEPIECE_AVAILABLE else '❌ No'}")
         gr.Markdown(f"**Dependencies Available**: {'✅ Yes' if DEPENDENCIES_AVAILABLE else '❌ No'}")
         gr.Markdown("**Architecture**: ✅ OpenLLM Custom GPTModel (From Uploaded Files)")
         
         # Main Content Area
+        # Two-column layout for configuration and status
         with gr.Row():
             
             # Left Column: Training Configuration
+            # Contains all the training hyperparameters and settings
             with gr.Column(scale=1):
                 gr.Markdown("## 📊 Training Configuration")
                 
                 # Model Size Selection
+                # Allows users to choose which base model to train from
                 model_size = gr.Dropdown(
                     choices=["small", "medium", "large"],
                     value="small",
-                    label="Model Size"
+                    label="Model Size",
+                    info="Select the base model size to train from"
                 )
                 
                 # Training Steps Configuration
+                # Controls the number of training iterations
                 max_steps = gr.Slider(
                     minimum=100,
                     maximum=10000,
                     value=1000,
                     step=100,
-                    label="Max Training Steps"
+                    label="Max Training Steps",
+                    info="Number of training iterations (100-10,000)"
                 )
                 
                 # Learning Rate Configuration
+                # Controls the learning rate for the optimizer
                 learning_rate = gr.Slider(
                     minimum=1e-5,
                     maximum=1e-3,
                     value=3e-4,
                     step=1e-5,
-                    label="Learning Rate"
+                    label="Learning Rate",
+                    info="Training rate (0.00001-0.001)"
                 )
                 
                 # Batch Size Configuration
+                # Controls the number of samples per training batch
                 batch_size = gr.Slider(
                     minimum=1,
                     maximum=16,
                     value=4,
                     step=1,
-                    label="Batch Size"
+                    label="Batch Size",
+                    info="Samples per training batch (1-16)"
                 )
             
             # Right Column: Training Status and Controls
+            # Contains status display and control buttons
             with gr.Column(scale=1):
                 gr.Markdown("## 🎯 Training Status")
                 
                 # Training Status Display
+                # Shows current training status and any error messages
                 status_text = gr.Textbox(
                     value="Ready to start training" if OPENLLM_AVAILABLE else "OpenLLM not available",
                     label="Current Status",
                     interactive=False,
-                    lines=5
+                    lines=5,
+                    info="Shows current training status and progress updates"
                 )
                 
                 # Progress Information
+                # Displays detailed training progress in JSON format
                 progress_info = gr.JSON(
                     value=trainer.get_training_progress(),
-                    label="Training Progress"
+                    label="Training Progress",
+                    info="Real-time training progress information"
                 )
                 
                 # Training Control Buttons
+                # Buttons to start and stop training
                 with gr.Row():
                     start_btn = gr.Button("🚀 Start Training", variant="primary")
                     stop_btn = gr.Button("⏹️ Stop Training", variant="stop")
         
         # Instructions Section
+        # Provides detailed instructions for using the training interface
         gr.Markdown("## 📋 OpenLLM Training Instructions")
         gr.Markdown("""
         This interface uses **OpenLLM's actual custom model architecture** from uploaded files:
@@ -562,6 +879,7 @@ def main():
         """)
         
         # Resource Links Section
+        # Provides links to related models and resources
         gr.Markdown("## 🔗 Model Resources")
         gr.Markdown("""
         - [📚 7k Small Model](https://huggingface.co/lemms/openllm-small-extended-7k)
@@ -571,15 +889,45 @@ def main():
         """)
         
         # Training Function Definition
+        # This function is called when the Start Training button is clicked
         def start_complete_training(model_size, max_steps, learning_rate, batch_size):
             """
             Execute the complete training process using OpenLLM's approach.
+            
+            This function orchestrates the entire training pipeline:
+            1. Validates OpenLLM availability
+            2. Creates training configuration
+            3. Loads model and tokenizer
+            4. Prepares dataset
+            5. Sets up training environment
+            6. Executes training
+            7. Saves and uploads the trained model
+            
+            The function provides comprehensive error handling and status updates
+            throughout the training process.
+            
+            Args:
+                model_size: Size of the model to train ("small", "medium", "large")
+                max_steps: Maximum number of training steps
+                learning_rate: Learning rate for the optimizer
+                batch_size: Batch size for training
+                
+            Returns:
+                Status message indicating the result of the training process
             """
+            # Validate OpenLLM availability
             if not OPENLLM_AVAILABLE:
                 return "❌ OpenLLM custom model architecture not available. Please check the installation."
             
             try:
+                print(f"🚀 Starting complete training process...")
+                print(f"   - Model size: {model_size}")
+                print(f"   - Max steps: {max_steps}")
+                print(f"   - Learning rate: {learning_rate}")
+                print(f"   - Batch size: {batch_size}")
+                
                 # Create training configuration
+                # This encapsulates all training parameters
                 config = TrainingConfig(
                     model_size=model_size,
                     max_steps=max_steps,
@@ -588,38 +936,55 @@ def main():
                 )
                 
                 # Step 1: Load model and tokenizer using OpenLLM's approach
+                print("🔄 Step 1: Loading model and tokenizer...")
                 status = trainer.load_model_and_tokenizer(model_size)
                 if "❌" in status:
                     return status
                 
                 # Step 2: Prepare dataset
+                print("🔄 Step 2: Preparing dataset...")
                 status = trainer.prepare_dataset()
                 if "❌" in status:
                     return status
                 
                 # Step 3: Setup training
+                print("🔄 Step 3: Setting up training...")
                 status = trainer.setup_training(config)
                 if "❌" in status:
                     return status
                 
                 # Step 4: Execute training
+                print("🔄 Step 4: Executing training...")
                 status = trainer.train_model(config)
                 if "❌" in status:
                     return status
                 
                 # Step 5: Save and upload model
+                print("🔄 Step 5: Saving and uploading model...")
                 status = trainer.save_and_upload_model(config)
                 
+                print("🎉 Complete training process finished!")
                 return f"🚀 Complete training process finished!\n{status}"
                 
             except Exception as e:
+                print(f"❌ Training process failed: {str(e)}")
                 return f"❌ Training process failed: {str(e)}"
         
         def update_progress():
-            """Update the progress display."""
+            """
+            Update the progress display.
+            
+            This function is called periodically to update the progress
+            information displayed in the Gradio interface. It returns the
+            current training progress from the trainer.
+            
+            Returns:
+                Current training progress dictionary
+            """
             return trainer.get_training_progress()
         
         # Connect UI Components to Functions
+        # This connects the Start Training button to the training function
         start_btn.click(
             fn=start_complete_training,
             inputs=[model_size, max_steps, learning_rate, batch_size],
@@ -627,9 +992,11 @@ def main():
         )
         
         # Auto-refresh progress every 5 seconds during training
+        # This ensures the progress display stays up to date
         demo.load(update_progress, outputs=[progress_info])
         
         # Application Footer
+        # Provides attribution and technical information
         gr.Markdown("---")
         gr.Markdown("**Author**: Louis Chua Bean Chong | **Project**: OpenLLM | **License**: GPL-3.0")
         gr.Markdown("**Architecture**: OpenLLM Custom GPTModel (From Uploaded Files)")
@@ -638,5 +1005,7 @@ def main():
     return demo
 
 if __name__ == "__main__":
+    # Launch the Gradio application
+    # This starts the web interface for the training application
     demo = main()
     demo.launch()