feat: Sync training infrastructure from main repository

app.py

@@ -1,223 +1,1024 @@
 #!/usr/bin/env python3
 """
-OpenLLM Training Space - Main Application
+OpenLLM Training Space Application - Fixed with Uploaded Modules
 
-This is the main entry point for the Hugging Face Space.
-It provides a web interface for running OpenLLM training with authentication.
+This version imports OpenLLM modules from the uploaded files in the HF Space:
+- Imports model.py and data_loader.py that were uploaded to the 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: GPLv3
+License: GPL-3.0
+Version: 2.1.1
+Last Updated: 2024
 """
 
-import os
-import sys
 import gradio as gr
+import torch
+import torch.nn as nn
+import os
+import time
+import math
+import gc
+from typing import Dict, Any, Optional
+import threading
+from dataclasses import dataclass
 from pathlib import Path
 
-# Import our authentication and training modules
+# 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 - 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 for the complete training pipeline
 try:
-    from space_auth_test import test_space_authentication
-    from openllm_training_with_auth import OpenLLMTrainingManager
-    MODULES_AVAILABLE = True
+    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:
-    MODULES_AVAILABLE = False
-    print(f"❌ Required modules not available: {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.
+    
+    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
+    batch_size: int
+    output_dir: str = "./openllm-trained"
+    save_steps: int = 100
+    logging_steps: int = 10
+    warmup_steps: int = 50
+    gradient_accumulation_steps: int = 4
 
-def create_space_interface():
-    """Create the Gradio interface for the Space."""
+class OpenLLMTrainer:
+    """
+    Complete training implementation using OpenLLM's actual architecture.
+    
+    This class handles the entire training pipeline including:
+    - Model loading using OpenLLM's custom GPTModel
+    - Tokenizer loading using sentencepiece.SentencePieceProcessor()
+    - Dataset preparation using OpenLLM's TextDataLoader
+    - Training execution using OpenLLM's approach
+    - 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 run_authentication_test():
-        """Run the authentication test and return results."""
+    def __init__(self):
+        """
+        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 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:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
+            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
             
-            # Capture output from authentication test
-            import io
-            import contextlib
+        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"
             
-            output = io.StringIO()
-            with contextlib.redirect_stdout(output):
-                success = test_space_authentication()
+            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")
             
-            result = output.getvalue()
+            # 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"   - 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)}"
             
-            if success:
-                return f"✅ Authentication Test Results:\n\n{result}"
-            else:
-                return f"❌ Authentication Test Failed:\n\n{result}"
+            # 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"  # Specific file name for OpenLLM
+                )
+                
+                print(f"✅ Tokenizer downloaded to: {tokenizer_path}")
+                print(f"   - Source: {model_name}")
+                print(f"   - File: tokenizer.model")
+                
+                # 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"   - Tokenizer type: {type(sp_processor).__name__}")
+                
+            except Exception as e:
+                print(f"❌ Failed to load tokenizer: {e}")
+                return f"❌ Failed to load OpenLLM tokenizer: {str(e)}"
+            
+            return f"✅ Successfully loaded OpenLLM {model_size} model with custom architecture"
+            
         except Exception as e:
-            return f"❌ Error running authentication test: {e}"
+            return f"❌ Failed to load OpenLLM model and tokenizer: {str(e)}"
     
-    def run_training(model_size, training_steps):
-        """Run the OpenLLM training with authentication."""
+    def prepare_dataset(self) -> str:
+        """
+        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:
-            if not MODULES_AVAILABLE:
-                return "❌ Required modules not available. Please check deployment."
+            # Verify dependencies are available
+            if not DEPENDENCIES_AVAILABLE:
+                return "❌ Required dependencies not available"
             
-            # Capture output from training
-            import io
-            import contextlib
+            print("🔄 Loading training dataset...")
+            print("   - Loading from Hugging Face Hub dataset")
+            print("   - Using OpenLLM's data preparation approach")
             
-            output = io.StringIO()
-            with contextlib.redirect_stdout(output):
-                training_manager = OpenLLMTrainingManager()
-                repo_id = training_manager.run_training(
-                    model_size=model_size,
-                    steps=int(training_steps)
+            # 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,  # Maximum sequence length for training
+                    batch_size=4,  # Will be overridden by training config
+                    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}")
+                return f"❌ Failed to create data loader: {str(e)}"
+            
+            return f"✅ Successfully prepared dataset with {len(dataset['train'])} samples"
+            
+        except Exception as e:
+            return f"❌ Failed to prepare dataset: {str(e)}"
+    
+    def setup_training(self, config: TrainingConfig) -> str:
+        """
+        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 containing all hyperparameters
+            
+        Returns:
+            Status message indicating success or failure
+            Success: "✅ Training setup completed successfully"
+            Failure: "❌ Failed to setup training: {error details}"
+        """
+        try:
+            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}")
             
-            result = output.getvalue()
+            # Set up optimizer (AdamW with weight decay)
+            # This follows OpenLLM's optimization strategy
+            print("🔄 Setting up AdamW optimizer with weight decay...")
             
-            return f"✅ Training Results:\n\n{result}\n\n🎉 Model available at: https://huggingface.co/{repo_id}"
+            # 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},      # 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),  # 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,  # 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  # 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]  # 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"
+            
         except Exception as e:
-            return f"❌ Error running training: {e}"
+            return f"❌ Failed to setup training: {str(e)}"
     
-    def check_space_environment():
-        """Check the Space environment and configuration."""
+    def train_model(self, config: TrainingConfig, progress_callback=None) -> str:
+        """
+        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 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:
-            # Check if we're in a Space
-            space_vars = ["SPACE_ID", "SPACE_HOST", "SPACE_REPO_ID"]
-            is_space = any(os.getenv(var) for var in space_vars)
+            # Set training state
+            self.is_training = True
+            self.training_progress["status"] = "Training"
+            self.training_progress["total_steps"] = config.max_steps
             
-            # Check HF_TOKEN
-            hf_token = os.getenv("HF_TOKEN")
+            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}")
             
-            result = "🔍 Space Environment Check:\n\n"
+            # Training loop using OpenLLM's approach
+            self.model.train()  # Set model to training mode
+            accumulated_loss = 0.0  # Track loss across accumulation steps
+            self.optimizer.zero_grad()  # Clear gradients
             
-            if is_space:
-                result += "✅ Running in Hugging Face Space environment\n"
-                for var in space_vars:
-                    value = os.getenv(var)
-                    if value:
-                        result += f"   - {var}: {value}\n"
-            else:
-                result += "ℹ️ Running in local environment\n"
+            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 - compute gradients
+                loss.backward()
+                
+                # Update weights every gradient_accumulation_steps
+                if (batch_idx + 1) % config.gradient_accumulation_steps == 0:
+                    # Clip gradients for training stability
+                    # This prevents exploding gradients
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+                    
+                    # 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 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 at specified intervals
+                    if step % config.logging_steps == 0:
+                        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 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 for the next accumulation cycle
+                    accumulated_loss = 0.0
+                    
+                    # Clean up memory periodically
+                    if step % 100 == 0:
+                        gc.collect()
+                        print(f"🧹 Memory cleanup at step {step}")
             
-            if hf_token:
-                result += f"✅ HF access token found: {hf_token[:8]}...{hf_token[-4:]}\n"
-                result += "   - Source: HF access token in Space settings\n"
-            else:
-                result += "❌ HF access token not found\n"
-                result += "   - Please set HF_TOKEN in Space settings with HF access token\n"
+            # 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
             
-            result += f"\n📁 Available modules: {'✅' if MODULES_AVAILABLE else '❌'}"
+            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 result
+            return f"✅ Training completed successfully! Final step: {step}"
             
         except Exception as e:
-            return f"❌ Error checking environment: {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
     
-    # Create the Gradio interface
-    with gr.Blocks(title="OpenLLM Training Space", theme=gr.themes.Soft()) as interface:
-        gr.Markdown("""
-        # 🚀 OpenLLM Training Space
+    def _save_checkpoint(self, output_dir: str, step: int, is_best: bool = False) -> None:
+        """
+        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.
         
-        Welcome to the OpenLLM Training Space! This Space provides a complete environment for training OpenLLM models with automatic Hugging Face authentication and model upload.
+        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,                                    # 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 if this is the best model
+            if is_best:
+                best_path = os.path.join(output_dir, "best_model.pt")
+                torch.save(checkpoint, best_path)
+                print(f"💾 Best model saved: {best_path}")
+            
+            print(f"💾 Checkpoint saved: {checkpoint_path}")
+            
+        except Exception as e:
+            print(f"❌ Failed to save checkpoint: {e}")
+    
+    def save_and_upload_model(self, config: TrainingConfig) -> str:
+        """
+        Save the trained model and upload it to Hugging Face Hub.
         
-                 ## 🔐 Authentication
-         
-         This Space uses HF access token for secure authentication. The HF_TOKEN is automatically available from your Space settings.
+        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
         
-        ## 📋 Available Actions
+        The uploaded model will be available at:
+        https://huggingface.co/lemms/openllm-{size}-extended-8k
         
-        1. **Environment Check**: Verify Space configuration and authentication
-        2. **Authentication Test**: Test Hugging Face authentication
-        3. **Run Training**: Start OpenLLM training with automatic upload
-        """)
+        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 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}")
+                
+                # Create the repository first if it doesn't exist
+                try:
+                    from huggingface_hub import create_repo
+                    create_repo(
+                        repo_id=repo_id,
+                        repo_type="model",
+                        exist_ok=True,
+                        private=False
+                    )
+                    print(f"✅ Repository {repo_id} ready for upload")
+                except Exception as create_error:
+                    print(f"⚠️ Repository creation warning: {create_error}")
+                    print("   Continuing with upload attempt...")
+                
+                # 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,
+                    repo_type="model",
+                    commit_message=f"Add trained OpenLLM {config.model_size} model (8k steps)"
+                )
+                
+                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:
+            print(f"❌ Failed to save/upload model: {e}")
+            return f"❌ Failed to save/upload model: {str(e)}"
+    
+    def get_training_progress(self) -> Dict[str, Any]:
+        """
+        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
         
-        with gr.Tab("🔍 Environment Check"):
-            gr.Markdown("Check the Space environment and configuration.")
-            env_check_btn = gr.Button("Check Environment", variant="primary")
-            env_output = gr.Textbox(label="Environment Status", lines=10, interactive=False)
-            env_check_btn.click(check_space_environment, outputs=env_output)
-        
-        with gr.Tab("🔐 Authentication Test"):
-            gr.Markdown("Test Hugging Face authentication using HF access token.")
-            auth_test_btn = gr.Button("Run Authentication Test", variant="primary")
-            auth_output = gr.Textbox(label="Authentication Results", lines=15, interactive=False)
-            auth_test_btn.click(run_authentication_test, outputs=auth_output)
-        
-        with gr.Tab("🚀 Run Training"):
-            gr.Markdown("""
-            Start OpenLLM training with automatic model upload.
-            
-            **Training Parameters:**
-            - **Model Size**: Choose the model size (small, medium, large)
-            - **Training Steps**: Number of training steps (default: 8000)
-            
-            **Expected Results:**
-            - Training will complete successfully
-            - Model will be uploaded to Hugging Face Hub
-            - Repository will be created with proper model files
-            """)
-            
-            with gr.Row():
+        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",
-                    info="Choose the model size for training"
+                    info="Select the base model size to train from"
                 )
-                training_steps = gr.Number(
-                    value=8000,
-                    label="Training Steps",
-                    info="Number of training steps",
-                    minimum=1000,
-                    maximum=50000
+                
+                # 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",
+                    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",
+                    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",
+                    info="Samples per training batch (1-16)"
                 )
             
-            train_btn = gr.Button("Start Training", variant="primary", size="lg")
-            train_output = gr.Textbox(label="Training Results", lines=20, interactive=False)
-            
-            train_btn.click(
-                run_training,
-                inputs=[model_size, training_steps],
-                outputs=train_output
-            )
+            # 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,
+                    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"
+                )
+                
+                # 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")
         
-        with gr.Tab("📚 Documentation"):
-            gr.Markdown("""
-            ## 📖 Available Documentation
-            
-            - **HUGGINGFACE_SPACE_SETUP_GUIDE.md**: Complete setup guide
-            - **SPACE_AUTHENTICATION_SUMMARY.md**: Authentication summary
-            - **SPACE_READY_SUMMARY.md**: Deployment summary
+        # 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:
+        
+        ### **Step 1: Configure Parameters**
+        - **Model Size**: Select the base model to train from (small, medium, large)
+        - **Max Steps**: Number of training iterations (100-10,000)
+        - **Learning Rate**: Training rate (0.00001-0.001)
+        - **Batch Size**: Samples per training batch (1-16)
+        
+        ### **Step 2: Start Training**
+        - Click "Start Training" to begin the actual training process
+        - Uses OpenLLM's custom GPTModel class from uploaded files
+        - Uses sentencepiece.SentencePieceProcessor() for tokenization
+        - Compatible with OpenLLM's actual implementation
+        
+        ### **Step 3: Monitor Progress**
+        - Watch the status updates and progress information
+        - Training may take several minutes depending on steps
+        - The final model will be uploaded to Hugging Face Hub
+        
+        ### **Step 4: Access Results**
+        - Trained models are automatically pushed to: `lemms/openllm-{size}-extended-8k`
+        - Check the model repository for your trained model
+        - Use the model for inference or further training
+        """)
+        
+        # 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)
+        - [🎯 8k Small Model](https://huggingface.co/lemms/openllm-small-extended-8k)
+        - [📊 Training Dataset](https://huggingface.co/datasets/lemms/openllm-training-data)
+        - [📖 Main Project](https://github.com/louischua/openllm)
+        """)
+        
+        # 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.
             
-            ## 🔧 Available Scripts
+            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
             
-            - **space_auth_test.py**: Authentication verification
-            - **openllm_training_with_auth.py**: Complete training script
-            - **integrate_auth_into_training.py**: Integration guide
-            - **setup_hf_space_auth.py**: Space authentication setup
-            - **verify_space_auth.py**: Space verification script
+            The function provides comprehensive error handling and status updates
+            throughout the training process.
             
-            ## 🎯 Quick Start
+            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."
             
-            1. Check the environment to verify configuration
-            2. Run authentication test to ensure GitHub secrets are working
-            3. Start training with your desired parameters
-            4. Monitor the training progress and model upload
+            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,
+                    learning_rate=learning_rate,
+                    batch_size=batch_size
+                )
+                
+                # 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.
             
-            ## 🔒 Security
+            This function is called periodically to update the progress
+            information displayed in the Gradio interface. It returns the
+            current training progress from the trainer.
             
-            - HF_TOKEN is securely stored in GitHub repository secrets
-            - No hardcoded tokens in any scripts
-            - Automatic cleanup of test repositories
-            - Proper error handling and logging
-            """)
+            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],
+            outputs=[status_text]
+        )
+        
+        # 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)")
+        gr.Markdown("**Tokenizer**: sentencepiece.SentencePieceProcessor()")
     
-    return interface
-
+    return demo
 
 if __name__ == "__main__":
-    # Create and launch the interface
-    interface = create_space_interface()
-    interface.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    # Launch the Gradio application
+    # This starts the web interface for the training application
+    demo = main()
+    demo.launch()

requirements.txt

@@ -1,26 +1,51 @@
-# OpenLLM Training Space Requirements
-# Core dependencies for Space deployment
+# Complete Training Dependencies for OpenLLM Space - Updated for Gradio 4.44.1
+# This file includes all necessary packages for real model training
 
-# Hugging Face Hub for authentication and model upload
-huggingface_hub>=0.19.0
+# Core Machine Learning Framework
+torch>=2.0.0                    # PyTorch deep learning framework
+torchvision>=0.15.0             # Computer vision utilities
+torchaudio>=2.0.0               # Audio processing utilities
 
-# Gradio for web interface (latest stable version with security fixes)
-gradio>=5.31.0
+# Hugging Face Ecosystem - Complete Training Stack
+transformers>=4.30.0            # Pre-trained models and training utilities
+datasets>=2.12.0                # Dataset loading and processing
+tokenizers>=0.13.0              # Fast tokenization library
+sentencepiece>=0.1.99           # SentencePiece tokenization (CRITICAL for OpenLLM models)
+huggingface_hub>=0.34.0         # Hugging Face Hub integration
+accelerate>=0.20.0              # Distributed training acceleration
 
-# PyTorch for model training
-torch>=2.0.0
-torchvision>=0.15.0
+# User Interface Framework - Updated to 4.44.1
+gradio==4.44.1                  # Web UI framework for ML applications (fixed version)
 
-# Transformers for model handling
-transformers>=4.35.0
+# Data Processing and Scientific Computing
+numpy>=1.24.0                   # Numerical computing library
+pandas>=2.0.0                   # Data manipulation and analysis
+scipy>=1.10.0                   # Scientific computing utilities
 
-# SentencePiece for tokenization
-sentencepiece>=0.1.99
+# Progress and Monitoring
+tqdm>=4.65.0                    # Progress bars for long-running operations
+psutil>=5.9.0                   # System and process utilities
 
-# NumPy and other utilities
-numpy>=1.24.0
-pandas>=2.0.0
+# Memory and Performance Optimization
+bitsandbytes>=0.41.0            # Quantization utilities for memory efficiency
+peft>=0.4.0                     # Parameter-Efficient Fine-Tuning
 
-# Additional utilities
-requests>=2.31.0
-tqdm>=4.65.0
+# Logging and Debugging
+wandb>=0.15.0                   # Experiment tracking (optional)
+tensorboard>=2.13.0             # Training visualization (optional)
+
+# Additional Utilities
+requests>=2.31.0                # HTTP library for API calls
+pillow>=9.5.0                   # Image processing (if needed)
+matplotlib>=3.7.0               # Plotting and visualization
+seaborn>=0.12.0                 # Statistical data visualization
+
+# Development and Testing (optional)
+pytest>=7.4.0                   # Testing framework
+black>=23.0.0                   # Code formatting
+flake8>=6.0.0                   # Code linting
+
+# Note: These versions are compatible with Hugging Face Spaces
+# and provide stable training performance for OpenLLM models
+# Gradio 4.44.1 fixes compatibility issues with JSON components
+# SentencePiece is CRITICAL for OpenLLM model tokenization

training/data_loader.py

@@ -113,7 +113,9 @@ class TextDataLoader:
 
         # Initialize data attribute for testing compatibility
         # Load a small sample of data for testing purposes
-        self.data = self._read_chunk(0, min(self.chunk_size, 100))  # Load up to 100 passages for testing
+        self.data = self._read_chunk(
+            0, min(self.chunk_size, 100)
+        )  # Load up to 100 passages for testing
 
         # Set random seed for reproducibility
         random.seed(seed)

training/model.py

@@ -514,7 +514,7 @@ class GPTModel(nn.Module):
         # Language modeling head
         # Always compute full logits for training and evaluation
         logits = self.lm_head(x)
-        
+
         if targets is not None:
             # If we have targets, compute loss
             loss = F.cross_entropy(