second

app.py

@@ -1,44 +1,51 @@
 import sys
 import logging
 from pathlib import Path
+import os
+import torch
 from transformers import set_seed
 
+# Set environment variables for memory optimization
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
 # Import the necessary modules from your project
-sys.path.append("./scripts")  # Adjust path if needed
+sys.path.append("./scripts")
 from scripts.models.model import NarrativeClassifier
 from scripts.models.dataset import NarrativeDataset
 from scripts.config.config import TrainingConfig
-from scripts.data_processing.data_preparation import AdvancedNarrativeProcessor
+from scripts.data_processing.advanced_preprocessor import AdvancedNarrativeProcessor
 from scripts.training.trainer import NarrativeTrainer
 
+def setup_logging():
+    """Setup logging configuration"""
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    return logging.getLogger(__name__)
+
 def main():
     # Set up logging
-    logging.basicConfig(level=logging.INFO)
-    logger = logging.getLogger(__name__)
+    logger = setup_logging()
     logger.info("Initializing training process...")
-    import os
-   
     
-    # Set up logging
-    logging.basicConfig(level=logging.INFO)
-    logger = logging.getLogger(__name__)
-    logger.info("Initializing training process...")
-    import os
-    import spacy
-    
-    # Download and load SpaCy model dynamically
-    try:
-        spacy.load("en_core_web_sm")
-    except OSError:
-        logger.info("Downloading SpaCy model 'en_core_web_sm'...")
-        os.system("python -m spacy download en_core_web_sm")
-
-    # Set a random seed for reproducibility
+    # Set random seeds for reproducibility
     set_seed(42)
+    torch.manual_seed(42)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(42)
+    
+    # Clear GPU cache if available
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        logger.info(f"CUDA available. Using GPU: {torch.cuda.get_device_name(0)}")
+        logger.info(f"Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
     
     # Load and process the dataset
-    annotations_file = "./data/subtask-2-annotations.txt"  # Adjust path as needed
-    raw_dir = "./data/raw"  # Adjust path as needed
+    annotations_file = "./data/subtask-2-annotations.txt"
+    raw_dir = "./data/raw"
     logger.info("Loading and processing dataset...")
     
     processor = AdvancedNarrativeProcessor(
@@ -47,41 +54,63 @@ def main():
     )
     processed_data = processor.load_and_process_data()
     
-    # Split processed data into training and validation sets
+    # Create datasets
     train_dataset = NarrativeDataset(processed_data['train'])
     val_dataset = NarrativeDataset(processed_data['val'])
     logger.info(f"Loaded dataset with {len(train_dataset)} training samples and {len(val_dataset)} validation samples.")
     
-    # Initialize the model
+    # Initialize model
     logger.info("Initializing the model...")
-    model = NarrativeClassifier(num_labels=train_dataset.get_num_labels())
+    model = NarrativeClassifier(
+        num_labels=train_dataset.get_num_labels(),
+        model_name="microsoft/deberta-v3-large"
+    )
     
-    # Define training configuration
+    # Define optimized training configuration
     config = TrainingConfig(
-        output_dir=Path("./output"),  # Save outputs in this directory
+        output_dir=Path("./output"),
         num_epochs=5,
-        batch_size=16,
+        batch_size=4,  # Reduced batch size for memory
         learning_rate=2e-5,
         warmup_ratio=0.1,
         weight_decay=0.01,
         max_grad_norm=1.0,
-        eval_steps=100,
-        save_steps=100
+        eval_steps=50,
+        save_steps=50,
+        fp16=True,  # Enable mixed precision
+        gradient_accumulation_steps=4,  # Gradient accumulation
+        max_length=256  # Reduced sequence length
     )
-    logger.info(f"Training configuration: {config}")
+    logger.info("Training configuration:")
+    for key, value in vars(config).items():
+        logger.info(f"  {key}: {value}")
     
-    # Initialize the trainer
-    trainer = NarrativeTrainer(
-        model=model,
-        train_dataset=train_dataset,
-        val_dataset=val_dataset,
-        config=config
-    )
-    
-    # Start the training process
-    logger.info("Starting the training process...")
-    trainer.train()
-    logger.info("Training completed successfully!")
+    try:
+        # Initialize trainer
+        trainer = NarrativeTrainer(
+            model=model,
+            train_dataset=train_dataset,
+            val_dataset=val_dataset,
+            config=config
+        )
+        
+        # Start training
+        logger.info("Starting the training process...")
+        history = trainer.train()
+        
+        # Log final metrics
+        logger.info("Training completed successfully!")
+        logger.info("Final metrics:")
+        logger.info(f"  Best validation F1: {trainer.best_val_f1:.4f}")
+        logger.info(f"  Final training loss: {history['train_loss'][-1]:.4f}")
+        
+    except Exception as e:
+        logger.error(f"Training failed with error: {str(e)}")
+        raise
+    finally:
+        # Clean up
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 if __name__ == "__main__":
-    main()
+    main()

requirements.txt

@@ -8,3 +8,4 @@ sentencepiece
 pandas
 numpy
 spacy
+accelerate

scripts/config/config.py

@@ -12,11 +12,14 @@ class TrainingConfig:
     
     # Training parameters
     num_epochs: int = 5
-    batch_size: int = 8
     learning_rate: float = 2e-5
     warmup_ratio: float = 0.1
     weight_decay: float = 0.01
     max_grad_norm: float = 1.0
+    gradient_accumulation_steps: int = 4
+    fp16: bool = True  # Enable mixed precision training
+    max_length: int = 256  # Reduce from 512
+    batch_size: int = 4  # Reduce from 8
     
     # Data parameters
     max_length: int = 512
@@ -45,15 +48,3 @@ if __name__ == "__main__":
     print(f"Learning rate: {default_config.learning_rate}")
     print(f"Device: {default_config.device}")
     
-    # Create custom config
-    custom_config = TrainingConfig(
-        batch_size=16,
-        num_epochs=10,
-        learning_rate=1e-5
-    )
-    
-    print("\n=== Custom Configuration ===")
-    print(f"Model name: {custom_config.model_name}")  # Uses default
-    print(f"Batch size: {custom_config.batch_size}")  # Customized
-    print(f"Learning rate: {custom_config.learning_rate}")  # Customized
-    print(f"Number of epochs: {custom_config.num_epochs}")  # Customized

scripts/training/trainer.py

@@ -8,12 +8,10 @@ import numpy as np
 from sklearn.metrics import f1_score, precision_score, recall_score
 import json
 from datetime import datetime
-
+from torch.cuda.amp import autocast, GradScaler
 
 class NarrativeTrainer:
-    """
-    Comprehensive trainer for narrative classification with GPU support.
-    """
+    """Comprehensive trainer for narrative classification with GPU memory optimizations"""
     def __init__(
         self,
         model,
@@ -21,29 +19,43 @@ class NarrativeTrainer:
         val_dataset,
         config,
     ):
+        # Setup basics
         self.setup_logging()
         self.logger = logging.getLogger(__name__)
         
-        # Set device
+        # Store config first
+        self.config = config
+        
+        # Setup device
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.logger.info(f"Using device: {self.device}")
         
+        # Clear GPU cache if using CUDA
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
         # Initialize model and components
         self.model = model.to(self.device)
         self.train_dataset = train_dataset
         self.val_dataset = val_dataset
-        self.config = config
         
+        # Initialize training state
         self.current_epoch = 0
         self.global_step = 0
         self.best_val_f1 = 0.0
         
+        # Initialize mixed precision training
+        self.scaler = GradScaler(enabled=self.config.fp16)
+        
+        # Setup training components
         self.setup_training()
         
+        # Setup output directory
         self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         self.output_dir = Path(config.output_dir) / self.timestamp
         self.output_dir.mkdir(parents=True, exist_ok=True)
         
+        # Save config and initialize history
         self.save_config()
         self.history = {
             'train_loss': [],
@@ -54,6 +66,7 @@ class NarrativeTrainer:
         }
     
     def setup_logging(self):
+        """Initialize logging configuration"""
         logging.basicConfig(
             level=logging.INFO,
             format='%(asctime)s - %(levelname)s - %(message)s',
@@ -61,27 +74,33 @@ class NarrativeTrainer:
         )
     
     def setup_training(self):
-        """Initialize dataloaders, optimizer, and scheduler."""
+        """Initialize training components with memory optimizations"""
+        # Create dataloaders
         self.train_loader = DataLoader(
             self.train_dataset,
             batch_size=self.config.batch_size,
             shuffle=True,
-            num_workers=4
+            num_workers=4,
+            pin_memory=True  # Optimize data transfer to GPU
         )
         
         self.val_loader = DataLoader(
             self.val_dataset,
             batch_size=self.config.batch_size,
-            num_workers=4
+            num_workers=4,
+            pin_memory=True
         )
         
+        # Setup optimizer
         self.optimizer = torch.optim.AdamW(
             self.model.parameters(),
             lr=self.config.learning_rate,
             weight_decay=self.config.weight_decay
         )
         
-        num_training_steps = len(self.train_loader) * self.config.num_epochs
+        # Setup scheduler with gradient accumulation steps
+        num_update_steps_per_epoch = len(self.train_loader) // self.config.gradient_accumulation_steps
+        num_training_steps = num_update_steps_per_epoch * self.config.num_epochs
         num_warmup_steps = int(num_training_steps * self.config.warmup_ratio)
         
         self.scheduler = get_linear_schedule_with_warmup(
@@ -93,66 +112,106 @@ class NarrativeTrainer:
         self.criterion = torch.nn.BCEWithLogitsLoss()
     
     def save_config(self):
-        """Save training configuration."""
+        """Save training configuration"""
         config_dict = {k: str(v) for k, v in vars(self.config).items()}
         config_path = self.output_dir / 'config.json'
         with open(config_path, 'w') as f:
             json.dump(config_dict, f, indent=4)
     
     def train_epoch(self):
-        """Train model for one epoch."""
+        """Train for one epoch with memory optimizations"""
         self.model.train()
         total_loss = 0
-        pbar = tqdm(self.train_loader, desc=f'Epoch {self.current_epoch + 1}/{self.config.num_epochs}')
+        self.optimizer.zero_grad()
         
-        for batch in pbar:
-            batch = {k: v.to(self.device) for k, v in batch.items()}
+        pbar = tqdm(enumerate(self.train_loader), 
+                   total=len(self.train_loader),
+                   desc=f'Epoch {self.current_epoch + 1}/{self.config.num_epochs}')
+        
+        for step, batch in pbar:
+            # Move batch to device
+            batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
             
-            self.optimizer.zero_grad()
-            outputs = self.model(
-                input_ids=batch['input_ids'],
-                attention_mask=batch['attention_mask'],
-                features=batch['features']
-            )
+            # Mixed precision forward pass
+            with autocast(enabled=self.config.fp16):
+                outputs = self.model(
+                    input_ids=batch['input_ids'],
+                    attention_mask=batch['attention_mask'],
+                    features=batch['features']
+                )
+                loss = self.criterion(outputs, batch['labels'])
+                loss = loss / self.config.gradient_accumulation_steps
             
-            loss = self.criterion(outputs, batch['labels'])
-            loss.backward()
+            # Scaled backward pass
+            self.scaler.scale(loss).backward()
             
-            torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config.max_grad_norm)
-            self.optimizer.step()
-            self.scheduler.step()
+            # Update weights if we've accumulated enough gradients
+            if (step + 1) % self.config.gradient_accumulation_steps == 0:
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(
+                    self.model.parameters(),
+                    self.config.max_grad_norm
+                )
+                
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.scheduler.step()
+                self.optimizer.zero_grad()
+            
+            # Update metrics
+            total_loss += loss.item() * self.config.gradient_accumulation_steps
+            avg_loss = total_loss / (step + 1)
+            pbar.set_postfix({'loss': f'{avg_loss:.4f}'})
             
-            total_loss += loss.item()
-            pbar.set_postfix({'loss': total_loss / (pbar.n + 1)})
             self.global_step += 1
             
+            # Evaluate if needed
             if self.global_step % self.config.eval_steps == 0:
                 self.evaluate()
+            
+            # Clear memory periodically
+            if step % 10 == 0:
+                torch.cuda.empty_cache()
+            
+            # Clear unnecessary tensors
+            del outputs
+            del loss
         
         return total_loss / len(self.train_loader)
     
     @torch.no_grad()
     def evaluate(self):
-        """Evaluate model performance."""
+        """Evaluate model with memory optimizations"""
         self.model.eval()
         total_loss = 0
         all_preds, all_labels = [], []
         
         for batch in tqdm(self.val_loader, desc="Evaluating"):
-            batch = {k: v.to(self.device) for k, v in batch.items()}
-            outputs = self.model(
-                input_ids=batch['input_ids'],
-                attention_mask=batch['attention_mask'],
-                features=batch['features']
-            )
+            batch = {k: v.to(self.device, non_blocking=True) for k, v in batch.items()}
+            
+            with autocast(enabled=self.config.fp16):
+                outputs = self.model(
+                    input_ids=batch['input_ids'],
+                    attention_mask=batch['attention_mask'],
+                    features=batch['features']
+                )
+                loss = self.criterion(outputs, batch['labels'])
             
-            loss = self.criterion(outputs, batch['labels'])
             total_loss += loss.item()
             
-            preds = torch.sigmoid(outputs) > 0.5
-            all_preds.append(preds.cpu().numpy())
-            all_labels.append(batch['labels'].cpu().numpy())
+            # CPU computations for predictions
+            preds = (torch.sigmoid(outputs) > 0.5).cpu().numpy()
+            labels = batch['labels'].cpu().numpy()
+            
+            all_preds.append(preds)
+            all_labels.append(labels)
+            
+            # Clear memory
+            del outputs
+            del loss
+            torch.cuda.empty_cache()
         
+        # Compute metrics
         all_preds = np.concatenate(all_preds, axis=0)
         all_labels = np.concatenate(all_labels, axis=0)
         
@@ -172,11 +231,13 @@ class NarrativeTrainer:
         return metrics
     
     def save_model(self, filename: str, metrics: dict = None):
+        """Save model checkpoint"""
         save_path = self.output_dir / filename
         torch.save({
             'model_state_dict': self.model.state_dict(),
             'optimizer_state_dict': self.optimizer.state_dict(),
             'scheduler_state_dict': self.scheduler.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict(),
             'epoch': self.current_epoch,
             'global_step': self.global_step,
             'best_val_f1': self.best_val_f1,
@@ -185,71 +246,34 @@ class NarrativeTrainer:
         self.logger.info(f"Model saved to {save_path}")
     
     def train(self):
-        """Run training for all epochs."""
+        """Run complete training loop"""
         self.logger.info("Starting training...")
-        for epoch in range(self.config.num_epochs):
-            self.current_epoch = epoch
-            train_loss = self.train_epoch()
-            self.history['train_loss'].append(train_loss)
+        try:
+            for epoch in range(self.config.num_epochs):
+                self.current_epoch = epoch
+                self.logger.info(f"Starting epoch {epoch + 1}/{self.config.num_epochs}")
+                
+                train_loss = self.train_epoch()
+                self.history['train_loss'].append(train_loss)
+                
+                val_metrics = self.evaluate()
+                self.history['val_loss'].append(val_metrics['loss'])
+                self.history['val_f1'].append(val_metrics['f1'])
+                self.history['val_precision'].append(val_metrics['precision'])
+                self.history['val_recall'].append(val_metrics['recall'])
+                
+                self.save_model(f'checkpoint_epoch_{epoch+1}.pt', val_metrics)
+                
+                # Save training history
+                history_path = self.output_dir / 'history.json'
+                with open(history_path, 'w') as f:
+                    json.dump(self.history, f, indent=4)
+                
+                self.logger.info(f"Epoch {epoch + 1} completed. Train loss: {train_loss:.4f}")
             
-            val_metrics = self.evaluate()
-            self.history['val_loss'].append(val_metrics['loss'])
-            self.history['val_f1'].append(val_metrics['f1'])
-            self.history['val_precision'].append(val_metrics['precision'])
-            self.history['val_recall'].append(val_metrics['recall'])
+            self.logger.info("Training completed successfully!")
+            return self.history
             
-            self.save_model(f'checkpoint_epoch_{epoch+1}.pt', val_metrics)
-            history_path = self.output_dir / 'history.json'
-            with open(history_path, 'w') as f:
-                json.dump(self.history, f, indent=4)
-        
-        self.logger.info("Training completed!")
-        return self.history
-
-
-if __name__ == "__main__":
-    import sys
-    sys.path.append("../../")
-    from scripts.models.model import NarrativeClassifier
-    from scripts.models.dataset import NarrativeDataset
-    from scripts.config.config import TrainingConfig
-    from scripts.data_processing.data_preparation import AdvancedNarrativeProcessor
-    
-    # Initialize training configuration
-    config = TrainingConfig(
-        output_dir=Path("./output"),
-        num_epochs=5,
-        batch_size=32,
-        learning_rate=5e-5,
-        weight_decay=0.01,
-        warmup_ratio=0.1,
-        max_grad_norm=1.0,
-        eval_steps=100
-    )
-    
-    # Load and process data
-    processor = AdvancedNarrativeProcessor(
-        annotations_file="../../data/subtask-2-annotations.txt",
-        raw_dir="../../data/raw"
-    )
-    processed_data = processor.load_and_process_data()
-    
-    # Create datasets
-    train_dataset = NarrativeDataset(processed_data['train'])
-    val_dataset = NarrativeDataset(processed_data['val'])
-    
-    # Initialize model
-    model = NarrativeClassifier(num_labels=train_dataset.get_num_labels())
-    
-    # Initialize trainer
-    trainer = NarrativeTrainer(
-        model=model,
-        train_dataset=train_dataset,
-        val_dataset=val_dataset,
-        config=config
-    )
-    
-    # Start full training
-    print("\n=== Starting Training ===")
-    trainer.train()
-    print("\nTraining completed successfully!")
+        except Exception as e:
+            self.logger.error(f"Training failed with error: {str(e)}")
+            raise