Update train.py

train.py

@@ -10,197 +10,389 @@ import argparse
 import time
 import math
 import glob
-from typing import Dict, List
+from typing import Dict, List, Optional
 from tqdm import tqdm
 import numpy as np
 import gc
+import logging
 from collections import defaultdict
 import multiprocessing
+
 # Import custom modules
 try:
     from model_slm import MixtureOfRecursions, count_parameters, TextGenerator
     from custom_tokenizer import TechnicalTokenizer
 except ImportError as e:
-    print(f"Import error: {e}")
-    exit(1)
+    raise ImportError(f"Failed to import custom modules: {e}")
+
+# Constants for configuration
+DEFAULT_MAX_LENGTH = 128
+DEFAULT_MAX_EXAMPLES = 50000
+DEFAULT_D_MODEL = 384
+DEFAULT_N_LAYERS = 6
+DEFAULT_N_HEADS = 6
+DEFAULT_EPOCHS = 15
+DEFAULT_BATCH_SIZE = 16
+DEFAULT_LEARNING_RATE = 5e-4
+DEFAULT_GRAD_ACCUM_STEPS = 1
+DEFAULT_EVAL_EVERY = 500
+DEFAULT_WARMUP_RATIO = 0.05
+DEFAULT_CHECKPOINT_DIR = "checkpoints"
+DEFAULT_LOG_LEVEL = "INFO"
+
+# Set up logging
+logging.basicConfig(
+    level=DEFAULT_LOG_LEVEL,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler("training.log")
+    ]
+)
+logger = logging.getLogger(__name__)
+
 class FastTechnicalTextDataset(Dataset):
-    """Ultra-fast dataset with aggressive optimizations for 4-5hr training"""    
-    def __init__(self, data_file: str, tokenizer: TechnicalTokenizer, max_length: int = 128, max_examples: int = 50000):
+    """Optimized dataset for fast loading and processing of technical text."""
+    
+    def __init__(
+        self,
+        data_file: str,
+        tokenizer: TechnicalTokenizer,
+        max_length: int = DEFAULT_MAX_LENGTH,
+        max_examples: int = DEFAULT_MAX_EXAMPLES
+    ):
+        """
+        Initialize the dataset with optimized loading.
+
+        Args:
+            data_file (str): Path to the training data file.
+            tokenizer (TechnicalTokenizer): Tokenizer for encoding text.
+            max_length (int): Maximum sequence length.
+            max_examples (int): Maximum number of examples to load.
+
+        Raises:
+            FileNotFoundError: If the data file does not exist.
+            ValueError: If max_length or max_examples is invalid.
+        """
+        if not os.path.exists(data_file):
+            raise FileNotFoundError(f"Data file not found: {data_file}")
+        if max_length <= 0 or max_examples <= 0:
+            raise ValueError("max_length and max_examples must be positive")
+
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.pad_token_id = tokenizer.vocab.get('<pad>', 0)
-        self.max_examples = max_examples        
-        print(f"FAST DATASET LOADING")
-        print(f"Data file: {data_file}")
-        print(f"Max sequence length: {max_length}")
-        print(f"Max examples: {max_examples}")        
-        start_time = time.time()
+        self.max_examples = max_examples
         self.examples = []
-        self._fast_load_data(data_file)        
-        load_time = time.time() - start_time
-        print(f" Loaded {len(self.examples)} examples in {load_time:.1f}s")        
+
+        logger.info(f"Loading dataset from {data_file} with max_length={max_length}, max_examples={max_examples}")
+        start_time = time.time()
+        self._fast_load_data(data_file)
         self._tensorize_data()
+        logger.info(f"Loaded {len(self.examples)} examples in {time.time() - start_time:.1f}s")
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         gc.collect()
-        torch.cuda.empty_cache() if torch.cuda.is_available() else None    
-    def _fast_load_data(self, data_file: str):
-        print("🔍 Fast reading file...")        
+
+    def _fast_load_data(self, data_file: str) -> None:
+        """Load and filter data efficiently."""
+        logger.info("Reading and filtering data...")
         with open(data_file, 'r', encoding='utf-8') as f:
-            lines = f.readlines()        
-        print(f"File has {len(lines)} lines")        
+            lines = f.readlines()
+
+        logger.info(f"File contains {len(lines)} lines")
         good_examples = []
-        seen_hashes = set()        
+        seen_hashes = set()
+
         for line in lines[:self.max_examples * 3]:
             line = line.strip()
-            if (50 <= len(line) <= 400 and
+            if (
+                50 <= len(line) <= 400 and
                 line.count(' ') >= 8 and
                 not line.lower().startswith(('http', 'www', 'ftp')) and
-                line.count('.') <= len(line) * 0.1):                
+                line.count('.') <= len(line) * 0.1
+            ):
                 line_hash = hash(line[:100])
                 if line_hash not in seen_hashes:
                     seen_hashes.add(line_hash)
                     good_examples.append(line)
                     if len(good_examples) >= self.max_examples:
-                        break        
-        print(f"After fast filtering: {len(good_examples)} quality examples")        
+                        break
+
+        logger.info(f"Filtered to {len(good_examples)} quality examples")
+
         batch_size = 1000
         for i in range(0, len(good_examples), batch_size):
-            batch = good_examples[i:i+batch_size]
+            batch = good_examples[i:i + batch_size]
             for line in batch:
                 try:
                     if not line.endswith('<|endoftext|>'):
-                        line += ' <|endoftext|>'                    
+                        line += ' <|endoftext|>'
                     tokens = self.tokenizer.encode_ids(line, add_special_tokens=True)
                     if 30 <= len(tokens) <= self.max_length:
                         if len(tokens) < self.max_length:
-                            tokens = tokens + [self.pad_token_id] * (self.max_length - len(tokens))
+                            tokens.extend([self.pad_token_id] * (self.max_length - len(tokens)))
                         self.examples.append(tokens)
-                except:
+                except Exception as e:
+                    logger.warning(f"Failed to process line: {e}")
                     continue
             if i % 5000 == 0:
-                print(f"Processed {len(self.examples)} examples...")        
-        print(f"Final dataset: {len(self.examples)} examples")    
-    def _tensorize_data(self):
-        print("Pre-tensorizing data for maximum speed...")
-        seq_len = self.max_length - 1        
+                logger.info(f"Processed {len(self.examples)} examples...")
+
+        logger.info(f"Final dataset size: {len(self.examples)} examples")
+
+    def _tensorize_data(self) -> None:
+        """Pre-tensorize data for faster training."""
+        logger.info("Pre-tensorizing data...")
+        seq_len = self.max_length - 1
         tensorized_examples = []
+
         for tokens in self.examples:
-            if len(tokens) < self.max_length:
-                continue            
+            if len(tokens) != self.max_length:
+                continue
             input_ids = torch.tensor(tokens[:-1], dtype=torch.long)
-            targets = torch.tensor(tokens[1:], dtype=torch.long)            
+            targets = torch.tensor(tokens[1:], dtype=torch.long)
             original_len = next((i for i, x in enumerate(tokens) if x == self.pad_token_id), self.max_length)
             mask_len = min(original_len, seq_len)
             attention_mask = torch.zeros(seq_len, dtype=torch.long)
-            attention_mask[:mask_len] = 1            
+            attention_mask[:mask_len] = 1
             tensorized_examples.append({
                 'input_ids': input_ids,
                 'targets': targets,
                 'attention_mask': attention_mask
             })
+
         self.examples = tensorized_examples
-        print("All data pre-tensorized")    
-    def __len__(self):
-        return len(self.examples)    
-    def __getitem__(self, idx):
+        logger.info("Data pre-tensorized successfully")
+
+    def __len__(self) -> int:
+        """Return the number of examples in the dataset."""
+        return len(self.examples)
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        """Return a single example from the dataset."""
         return self.examples[idx]
+
 class FastCosineScheduler:
-    def __init__(self, optimizer, total_steps: int, warmup_ratio: float = 0.05):
+    """Cosine learning rate scheduler with warmup."""
+    
+    def __init__(self, optimizer: optim.Optimizer, total_steps: int, warmup_ratio: float = DEFAULT_WARMUP_RATIO):
+        """
+        Initialize the cosine scheduler.
+
+        Args:
+            optimizer (optim.Optimizer): Optimizer to schedule.
+            total_steps (int): Total training steps.
+            warmup_ratio (float): Ratio of steps for warmup phase.
+
+        Raises:
+            ValueError: If total_steps or warmup_ratio is invalid.
+        """
+        if total_steps <= 0 or not 0 <= warmup_ratio <= 1:
+            raise ValueError("total_steps must be positive and warmup_ratio must be in [0, 1]")
+
         self.optimizer = optimizer
         self.total_steps = total_steps
         self.warmup_steps = int(total_steps * warmup_ratio)
         self.base_lr = optimizer.param_groups[0]['lr']
-        self.step_count = 0    
-    def step(self):
+        self.step_count = 0
+
+    def step(self) -> float:
+        """
+        Update the learning rate.
+
+        Returns:
+            float: Current learning rate.
+        """
         self.step_count += 1
         if self.step_count <= self.warmup_steps:
             lr = self.base_lr * self.step_count / self.warmup_steps
         else:
             progress = (self.step_count - self.warmup_steps) / (self.total_steps - self.warmup_steps)
             lr = self.base_lr * 0.5 * (1 + math.cos(math.pi * progress))
+        
         for param_group in self.optimizer.param_groups:
             param_group['lr'] = lr
         return lr
+
 class UltraFastTrainer:
-    def __init__(self, model, tokenizer, train_dataset, val_dataset=None, config=None):
+    """Trainer optimized for fast training of transformer models."""
+    
+    def __init__(
+        self,
+        model: nn.Module,
+        tokenizer: TechnicalTokenizer,
+        train_dataset: FastTechnicalTextDataset,
+        val_dataset: Optional[FastTechnicalTextDataset] = None,
+        config: Optional[Dict] = None
+    ):
+        """
+        Initialize the trainer.
+
+        Args:
+            model (nn.Module): The transformer model to train.
+            tokenizer (TechnicalTokenizer): Tokenizer for encoding/decoding.
+            train_dataset (FastTechnicalTextDataset): Training dataset.
+            val_dataset (Optional[FastTechnicalTextDataset]): Validation dataset.
+            config (Optional[Dict]): Training configuration.
+
+        Raises:
+            ValueError: If config contains invalid parameters.
+        """
         self.model = model
         self.tokenizer = tokenizer
         self.train_dataset = train_dataset
         self.val_dataset = val_dataset
-        self.config = config or {}        
+        self.config = config or {}
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        self.model.to(self.device)        
+        self.model.to(self.device)
+        
+        self._validate_config()
         self._fast_init_weights()
-        self._setup_fast_optimizer()        
-        epochs = self.config.get('epochs', 15)
-        batch_size = self.config.get('batch_size', 16)
+        self._setup_fast_optimizer()
+        
+        epochs = self.config.get('epochs', DEFAULT_EPOCHS)
+        batch_size = self.config.get('batch_size', DEFAULT_BATCH_SIZE)
         total_steps = len(train_dataset) // batch_size * epochs
         self.scheduler = FastCosineScheduler(self.optimizer, total_steps)
-        self.scaler = GradScaler()        
+        self.scaler = GradScaler() if self.device.type == 'cuda' else None
         self.global_step = 0
         self.best_loss = float('inf')
-        self.grad_accum_steps = self.config.get('gradient_accumulation_steps', 1)
-        self.eval_every = self.config.get('eval_every', 500)    
-    def _fast_init_weights(self):
-        def fast_init(module):
+        self.grad_accum_steps = self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS)
+        self.eval_every = self.config.get('eval_every', DEFAULT_EVAL_EVERY)
+
+    def _validate_config(self) -> None:
+        """Validate training configuration."""
+        if self.config.get('batch_size', DEFAULT_BATCH_SIZE) <= 0:
+            raise ValueError("batch_size must be positive")
+        if self.config.get('epochs', DEFAULT_EPOCHS) <= 0:
+            raise ValueError("epochs must be positive")
+        if self.config.get('learning_rate', DEFAULT_LEARNING_RATE) <= 0:
+            raise ValueError("learning_rate must be positive")
+        if self.config.get('gradient_accumulation_steps', DEFAULT_GRAD_ACCUM_STEPS) <= 0:
+            raise ValueError("gradient_accumulation_steps must be positive")
+
+    def _fast_init_weights(self) -> None:
+        """Initialize model weights."""
+        def fast_init(module: nn.Module) -> None:
             if isinstance(module, nn.Linear):
                 nn.init.normal_(module.weight, std=0.02)
                 if module.bias is not None:
                     nn.init.zeros_(module.bias)
             elif isinstance(module, nn.Embedding):
                 nn.init.normal_(module.weight, std=0.02)
-        self.model.apply(fast_init)    
-    def _setup_fast_optimizer(self):
-        lr = self.config.get('learning_rate', 5e-4)
+        self.model.apply(fast_init)
+        logger.info("Model weights initialized")
+
+    def _setup_fast_optimizer(self) -> None:
+        """Set up AdamW optimizer."""
+        lr = self.config.get('learning_rate', DEFAULT_LEARNING_RATE)
         params = [p for p in self.model.parameters() if p.requires_grad]
-        self.optimizer = optim.AdamW(params, lr=lr, betas=(0.9, 0.99), weight_decay=0.01, eps=1e-6)    
-    def compute_fast_loss(self, logits, targets, mask):
+        self.optimizer = optim.AdamW(params, lr=lr, betas=(0.9, 0.99), weight_decay=0.01, eps=1e-6)
+        logger.info(f"Optimizer initialized with learning rate: {lr}")
+
+    def compute_fast_loss(self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Compute masked cross-entropy loss.
+
+        Args:
+            logits (torch.Tensor): Model output logits of shape (batch_size, seq_len, vocab_size).
+            targets (torch.Tensor): Target token IDs of shape (batch_size, seq_len).
+            mask (torch.Tensor): Attention mask of shape (batch_size, seq_len).
+
+        Returns:
+            torch.Tensor: Computed loss.
+        """
         logits_flat = logits.view(-1, logits.size(-1))
         targets_flat = targets.view(-1)
         mask_flat = mask.view(-1).bool()
+        
         if not mask_flat.any():
             return torch.tensor(0.0, device=logits.device, requires_grad=True)
-        loss = F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])
-        return loss    
+        
+        return F.cross_entropy(logits_flat[mask_flat], targets_flat[mask_flat])
+
     def train_epoch_fast(self, epoch: int, dataloader: DataLoader) -> Dict[str, float]:
+        """
+        Train for one epoch.
+
+        Args:
+            epoch (int): Current epoch number.
+            dataloader (DataLoader): Training data loader.
+
+        Returns:
+            Dict[str, float]: Training metrics (loss, perplexity, epoch_time_min).
+        """
         self.model.train()
         total_loss = 0
         num_batches = 0
-        start_time = time.time()        
+        start_time = time.time()
+        
         progress_bar = tqdm(dataloader, desc=f"Epoch {epoch}", leave=False, miniters=50)
         for batch_idx, batch in enumerate(progress_bar):
             input_ids = batch['input_ids'].to(self.device, non_blocking=True)
             targets = batch['targets'].to(self.device, non_blocking=True)
-            mask = batch['attention_mask'].to(self.device, non_blocking=True)            
-            with autocast():
+            mask = batch['attention_mask'].to(self.device, non_blocking=True)
+            
+            with autocast(enabled=self.device.type == 'cuda'):
                 logits, comp_loss = self.model(input_ids, mask)
                 lm_loss = self.compute_fast_loss(logits, targets, mask)
                 total_loss_step = lm_loss + 0.0001 * comp_loss
                 if self.grad_accum_steps > 1:
-                    total_loss_step = total_loss_step / self.grad_accum_steps            
-            self.scaler.scale(total_loss_step).backward()
-            if (batch_idx + 1) % self.grad_accum_steps == 0:
-                self.scaler.unscale_(self.optimizer)
-                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
-                self.scaler.step(self.optimizer)
-                self.scaler.update()
-                self.optimizer.zero_grad(set_to_none=True)
-                self.scheduler.step()
-                self.global_step += 1            
+                    total_loss_step = total_loss_step / self.grad_accum_steps
+            
+            if self.scaler:
+                self.scaler.scale(total_loss_step).backward()
+                if (batch_idx + 1) % self.grad_accum_steps == 0:
+                    self.scaler.unscale_(self.optimizer)
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                    self.optimizer.zero_grad(set_to_none=True)
+                    self.scheduler.step()
+                    self.global_step += 1
+            else:
+                total_loss_step.backward()
+                if (batch_idx + 1) % self.grad_accum_steps == 0:
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
+                    self.optimizer.step()
+                    self.optimizer.zero_grad(set_to_none=True)
+                    self.scheduler.step()
+                    self.global_step += 1
+            
             total_loss += lm_loss.item()
             num_batches += 1
+            
             if batch_idx % 100 == 0:
                 current_loss = total_loss / num_batches
                 progress_bar.set_postfix({'loss': f"{current_loss:.3f}", 'ppl': f"{math.exp(min(current_loss, 10)):.1f}"})
-            if batch_idx % 200 == 0 and batch_idx > 0:
-                torch.cuda.empty_cache()       
-        epoch_time = time.time() - start_time
+            
+            if batch_idx % 200 == 0 and batch_idx > 0 and self.device.type == 'cuda':
+                torch.cuda.empty_cache()
+        
         avg_loss = total_loss / max(num_batches, 1)
-        return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10)), 'epoch_time_min': epoch_time / 60}    
+        return {
+            'loss': avg_loss,
+            'perplexity': math.exp(min(avg_loss, 10)),
+            'epoch_time_min': (time.time() - start_time) / 60
+        }
+
     def validate_fast(self, dataloader: DataLoader) -> Dict[str, float]:
+        """
+        Validate the model on the validation dataset.
+
+        Args:
+            dataloader (DataLoader): Validation data loader.
+
+        Returns:
+            Dict[str, float]: Validation metrics (loss, perplexity).
+        """
         self.model.eval()
         total_loss = 0
         num_batches = 0
-        max_val_batches = min(100, len(dataloader))        
+        max_val_batches = min(100, len(dataloader))
+        
         with torch.no_grad():
             for batch_idx, batch in enumerate(dataloader):
                 if batch_idx >= max_val_batches:
@@ -208,16 +400,32 @@ class UltraFastTrainer:
                 input_ids = batch['input_ids'].to(self.device, non_blocking=True)
                 targets = batch['targets'].to(self.device, non_blocking=True)
                 mask = batch['attention_mask'].to(self.device, non_blocking=True)
-                with autocast():
+                
+                with autocast(enabled=self.device.type == 'cuda'):
                     logits, _ = self.model(input_ids, mask)
                     loss = self.compute_fast_loss(logits, targets, mask)
+                
                 total_loss += loss.item()
-                num_batches += 1        
+                num_batches += 1
+        
         avg_loss = total_loss / max(num_batches, 1)
-        return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10))}    
-    def save_checkpoint_fast(self, epoch: int, metrics: Dict, save_dir: str = "checkpoints"):
+        return {'loss': avg_loss, 'perplexity': math.exp(min(avg_loss, 10))}
+
+    def save_checkpoint_fast(self, epoch: int, metrics: Dict, save_dir: str = DEFAULT_CHECKPOINT_DIR) -> Optional[str]:
+        """
+        Save a checkpoint if the loss improves.
+
+        Args:
+            epoch (int): Current epoch number.
+            metrics (Dict): Training and validation metrics.
+            save_dir (str): Directory to save checkpoints.
+
+        Returns:
+            Optional[str]: Path to the saved checkpoint or None.
+        """
         os.makedirs(save_dir, exist_ok=True)
         val_loss = metrics.get('val_loss', metrics.get('loss', float('inf')))
+        
         if val_loss < self.best_loss:
             self.best_loss = val_loss
             checkpoint = {
@@ -225,159 +433,211 @@ class UltraFastTrainer:
                 'model_state_dict': self.model.state_dict(),
                 'optimizer_state_dict': self.optimizer.state_dict(),
                 'metrics': metrics,
-                'scaler_state_dict': self.scaler.state_dict()
+                'scaler_state_dict': self.scaler.state_dict() if self.scaler else None
             }
             best_path = os.path.join(save_dir, "best_model.pt")
             torch.save(checkpoint, best_path)
-            print(f"New best! Loss: {val_loss:.4f}")
+            logger.info(f"New best checkpoint saved: {best_path}, Loss: {val_loss:.4f}")
             return best_path
-        return None    
-    def train_ultra_fast(self, num_epochs: int = 15, batch_size: int = 16):
-        print(f"\n ULTRA-FAST TRAINING")
-        print(f" Target: Loss < 2.0, PPL < 12")
-        print(f" Time target: 4-5 hours")
-        print(f" Epochs: {num_epochs}")
-        print(f" Batch size: {batch_size}")
-        print("-" * 60)        
+        return None
+
+    def train_ultra_fast(self, num_epochs: int = DEFAULT_EPOCHS, batch_size: int = DEFAULT_BATCH_SIZE) -> List[Dict]:
+        """
+        Train the model with optimized settings.
+
+        Args:
+            num_epochs (int): Number of training epochs.
+            batch_size (int): Batch size for training.
+
+        Returns:
+            List[Dict]: Training history with metrics for each epoch.
+        """
+        logger.info(f"Starting ultra-fast training: {num_epochs} epochs, batch_size={batch_size}")
+        logger.info("Target: Loss < 2.0, PPL < 12, Time: 4-5 hours")
+
         train_loader = DataLoader(
             self.train_dataset,
             batch_size=batch_size,
             shuffle=True,
-            num_workers=4,
-            pin_memory=True,
+            num_workers=min(multiprocessing.cpu_count(), 4),
+            pin_memory=self.device.type == 'cuda',
             persistent_workers=True,
             drop_last=True
-        )        
+        )
+
         val_loader = None
         if self.val_dataset:
             val_loader = DataLoader(
                 self.val_dataset,
                 batch_size=batch_size * 2,
                 shuffle=False,
-                num_workers=2,
-                pin_memory=True
-            )        
+                num_workers=min(multiprocessing.cpu_count() // 2, 2),
+                pin_memory=self.device.type == 'cuda'
+            )
+
         total_start_time = time.time()
-        history = []        
+        history = []
+
         for epoch in range(1, num_epochs + 1):
-            epoch_start = time.time()
-            print(f"\n EPOCH {epoch}/{num_epochs}")
-            train_metrics = self.train_epoch_fast(epoch, train_loader)            
+            logger.info(f"Starting epoch {epoch}/{num_epochs}")
+            train_metrics = self.train_epoch_fast(epoch, train_loader)
+            
             val_metrics = {}
             if val_loader and (epoch % 2 == 0 or epoch == num_epochs):
-                val_metrics = self.validate_fast(val_loader)            
-            epoch_time = time.time() - epoch_start
+                val_metrics = self.validate_fast(val_loader)
+            
+            epoch_time = train_metrics['epoch_time_min'] * 60
             epoch_info = {
                 'epoch': epoch,
                 'train_loss': train_metrics['loss'],
                 'train_ppl': train_metrics['perplexity'],
-                'epoch_time_min': epoch_time / 60
+                'epoch_time_min': train_metrics['epoch_time_min']
             }
             if val_metrics:
                 epoch_info.update({'val_loss': val_metrics['loss'], 'val_ppl': val_metrics['perplexity']})
-            history.append(epoch_info)            
+            
+            history.append(epoch_info)
+            
             elapsed_hours = (time.time() - total_start_time) / 3600
-            remaining_hours = elapsed_hours * (num_epochs - epoch) / epoch            
-            print(f"\n EPOCH {epoch} RESULTS:")
-            print(f" Epoch time: {epoch_time/60:.1f} min")
-            print(f" Total elapsed: {elapsed_hours:.1f}h")
-            print(f" Est. remaining: {remaining_hours:.1f}h")
-            print(f" Train Loss: {train_metrics['loss']:.4f}")
-            print(f" Train PPL: {train_metrics['perplexity']:.1f}")
+            remaining_hours = elapsed_hours * (num_epochs - epoch) / max(epoch, 1)
+            
+            logger.info(f"Epoch {epoch} results:")
+            logger.info(f"  Epoch time: {epoch_time/60:.1f} min")
+            logger.info(f"  Total elapsed: {elapsed_hours:.1f}h")
+            logger.info(f"  Est. remaining: {remaining_hours:.1f}h")
+            logger.info(f"  Train Loss: {train_metrics['loss']:.4f}")
+            logger.info(f"  Train PPL: {train_metrics['perplexity']:.1f}")
             if val_metrics:
-                print(f" Val Loss: {val_metrics['loss']:.4f}")
-                print(f" Val PPL: {val_metrics['perplexity']:.1f}")            
+                logger.info(f"  Val Loss: {val_metrics['loss']:.4f}")
+                logger.info(f"  Val PPL: {val_metrics['perplexity']:.1f}")
+            
             current_loss = val_metrics.get('loss', train_metrics['loss'])
             current_ppl = val_metrics.get('perplexity', train_metrics['perplexity'])
             if current_loss < 2.0 and current_ppl < 12:
-                print(f" TARGETS ACHIEVED!")
-                print(f" Loss: {current_loss:.4f} < 2.0")
-                print(f" PPL: {current_ppl:.1f} < 12")            
+                logger.info(f"Targets achieved: Loss={current_loss:.4f} < 2.0, PPL={current_ppl:.1f} < 12")
+            
             combined_metrics = {**train_metrics}
             if val_metrics:
                 combined_metrics.update({f"val_{k}": v for k, v in val_metrics.items()})
-            self.save_checkpoint_fast(epoch, combined_metrics)           
-            torch.cuda.empty_cache()
-            gc.collect()            
+            self.save_checkpoint_fast(epoch, combined_metrics)
+            
+            if self.device.type == 'cuda':
+                torch.cuda.empty_cache()
+            gc.collect()
+            
             if current_loss < 1.8 and current_ppl < 10:
-                print(f"EARLY STOPPING - Excellent performance achieved!")
-                break        
-        total_time = time.time() - total_start_time
-        print(f"\n TRAINING COMPLETED!")
-        print(f"Total time: {total_time/3600:.1f} hours")
-        print(f" Best loss: {self.best_loss:.4f}")
+                logger.info("Early stopping: Excellent performance achieved!")
+                break
+        
+        total_time = (time.time() - total_start_time) / 3600
+        logger.info(f"Training completed in {total_time:.1f} hours")
+        logger.info(f"Best loss: {self.best_loss:.4f}")
         return history
-def run_ultra_fast_training():
-    parser = argparse.ArgumentParser(description="Ultra-Fast Training for 4-5 Hours")
-    parser.add_argument("--train_file", default=None)
-    parser.add_argument("--val_file", default=None)
-    parser.add_argument("--tokenizer_dir", default="tokenizer")
-    parser.add_argument("--max_examples", type=int, default=50000)
-    parser.add_argument("--d_model", type=int, default=384)
-    parser.add_argument("--n_layers", type=int, default=6)
-    parser.add_argument("--n_heads", type=int, default=6)
-    parser.add_argument("--max_seq_len", type=int, default=128)
-    parser.add_argument("--epochs", type=int, default=15)
-    parser.add_argument("--batch_size", type=int, default=16)
-    parser.add_argument("--learning_rate", type=float, default=5e-4)
-    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
-    parser.add_argument("--eval_every", type=int, default=500)    
-    args = parser.parse_args()    
+
+def run_ultra_fast_training() -> int:
+    """
+    Run the ultra-fast training pipeline.
+
+    Returns:
+        int: Exit code (0 for success, 1 for failure).
+    """
+    parser = argparse.ArgumentParser(description="Ultra-Fast Training for MixtureOfRecursions Model")
+    parser.add_argument("--train_file", default=None, help="Path to training data file")
+    parser.add_argument("--val_file", default=None, help="Path to validation data file")
+    parser.add_argument("--tokenizer_dir", default="tokenizer", help="Directory for tokenizer files")
+    parser.add_argument("--max_examples", type=int, default=DEFAULT_MAX_EXAMPLES, help="Maximum number of training examples")
+    parser.add_argument("--d_model", type=int, default=DEFAULT_D_MODEL, help="Model embedding dimension")
+    parser.add_argument("--n_layers", type=int, default=DEFAULT_N_LAYERS, help="Number of transformer layers")
+    parser.add_argument("--n_heads", type=int, default=DEFAULT_N_HEADS, help="Number of attention heads")
+    parser.add_argument("--max_seq_len", type=int, default=DEFAULT_MAX_LENGTH, help="Maximum sequence length")
+    parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS, help="Number of training epochs")
+    parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for training")
+    parser.add_argument("--learning_rate", type=float, default=DEFAULT_LEARNING_RATE, help="Learning rate")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=DEFAULT_GRAD_ACCUM_STEPS, help="Gradient accumulation steps")
+    parser.add_argument("--eval_every", type=int, default=DEFAULT_EVAL_EVERY, help="Evaluate every N steps")
+    
+    args = parser.parse_args()
+    
     torch.manual_seed(42)
-    np.random.seed(42)    
-    print("Training My Model")
-    print("-" * 50)    
+    np.random.seed(42)
+    
+    logger.info("Starting ultra-fast training pipeline")
+    
     if args.train_file is None:
         patterns = ["*train*.txt", "*_train.txt"]
         files = []
         for pattern in patterns:
             files.extend(glob.glob(pattern))
-            files.extend(glob.glob(f"split_data/{pattern}"))
-            files.extend(glob.glob(f"data/{pattern}"))
+            files.extend(glob.glob(os.path.join("split_data", pattern)))
+            files.extend(glob.glob(os.path.join("data", pattern)))
         if files:
             args.train_file = files[0]
-            print(f"Found: {args.train_file}")
+            logger.info(f"Found training file: {args.train_file}")
         else:
-            print(" No training files found!")
-            return 1    
-    tokenizer = TechnicalTokenizer()
+            logger.error("No training files found!")
+            return 1
+    
     try:
+        tokenizer = TechnicalTokenizer()
         tokenizer.load(args.tokenizer_dir)
-        print(f"Tokenizer loaded. Vocab size: {tokenizer.get_vocab_size()}")
+        logger.info(f"Tokenizer loaded with vocab size: {tokenizer.get_vocab_size()}")
     except Exception as e:
-        print(f" Tokenizer error: {e}")
-        return 1    
-    print(" Creating ultra-fast dataset...")
-    train_dataset = FastTechnicalTextDataset(
-        args.train_file, tokenizer, args.max_seq_len, args.max_examples
-    )    
+        logger.error(f"Failed to load tokenizer: {e}")
+        return 1
+    
+    logger.info("Creating training dataset...")
+    try:
+        train_dataset = FastTechnicalTextDataset(
+            args.train_file, tokenizer, args.max_seq_len, args.max_examples
+        )
+    except Exception as e:
+        logger.error(f"Failed to create training dataset: {e}")
+        return 1
+    
     val_dataset = None
     if args.val_file and os.path.exists(args.val_file):
-        val_dataset = FastTechnicalTextDataset(
-            args.val_file, tokenizer, args.max_seq_len, max_examples=5000
-        )    
-    model = MixtureOfRecursions(
-        vocab_size=tokenizer.get_vocab_size(),
-        d_model=args.d_model,
-        n_layers=args.n_layers,
-        n_heads=args.n_heads,
-        max_seq_len=args.max_seq_len - 1, # Pass the actual sequence length to the model
-        padding_idx=tokenizer.vocab.get('<pad>', 0)
-    )    
+        try:
+            val_dataset = FastTechnicalTextDataset(
+                args.val_file, tokenizer, args.max_seq_len, max_examples=5000
+            )
+            logger.info("Validation dataset created")
+        except Exception as e:
+            logger.warning(f"Failed to create validation dataset: {e}")
+    
+    try:
+        model = MixtureOfRecursions(
+            vocab_size=tokenizer.get_vocab_size(),
+            d_model=args.d_model,
+            n_layers=args.n_layers,
+            n_heads=args.n_heads,
+            max_seq_len=args.max_seq_len - 1,
+            padding_idx=tokenizer.vocab.get('<pad>', 0)
+        )
+        logger.info("Model initialized")
+    except Exception as e:
+        logger.error(f"Failed to initialize model: {e}")
+        return 1
+    
     config = {
         'learning_rate': args.learning_rate,
         'gradient_accumulation_steps': args.gradient_accumulation_steps,
         'eval_every': args.eval_every,
         'batch_size': args.batch_size,
         'epochs': args.epochs
-    }    
-    trainer = UltraFastTrainer(model, tokenizer, train_dataset, val_dataset, config)
-    print(f"\n START TRAINING")
-    results = trainer.train_ultra_fast(args.epochs, args.batch_size)    
-    with open('ultra_fast_results.json', 'w') as f:
-        json.dump(results, f, indent=2)    
-    print("\n Training Completed!")
-    print(" Results saved to: ultra_fast_results.json")
-    return 0
+    }
+    
+    try:
+        trainer = UltraFastTrainer(model, tokenizer, train_dataset, val_dataset, config)
+        results = trainer.train_ultra_fast(args.epochs, args.batch_size)
+        
+        with open('ultra_fast_results.json', 'w') as f:
+            json.dump(results, f, indent=2)
+        logger.info("Training results saved to ultra_fast_results.json")
+        
+        return 0
+    except Exception as e:
+        logger.error(f"Training failed: {e}")
+        return 1
+
 if __name__ == "__main__":
     exit(run_ultra_fast_training())