Update train.py

train.py

@@ -1,17 +1,18 @@
 """
 Training module for Talmud language classifier
 Adapted from talmud_language_classifier.py for Hugging Face Spaces integration
+Optimized for class imbalance and better performance
 """
 
 import copy
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
 from collections import Counter
 from sklearn.model_selection import train_test_split, KFold
 from sklearn.preprocessing import LabelEncoder
-from sklearn.metrics import f1_score
+from sklearn.metrics import f1_score, classification_report
 import numpy as np
 import io
 import os
@@ -21,10 +22,14 @@ import pickle
 MAX_LEN = 100
 VOCAB_SIZE = 10000
 EMBEDDING_DIM = 128
-HIDDEN_DIM = 128
-NUM_EPOCHS = 10  # Epochs per fold
+HIDDEN_DIM = 256  # Increased for better capacity
+NUM_EPOCHS = 30  # Increased epochs with early stopping
 BATCH_SIZE = 16
 N_SPLITS = 5  # Number of folds for cross-validation
+EARLY_STOPPING_PATIENCE = 5  # Stop if no improvement for 5 epochs
+LEARNING_RATE = 0.001
+WEIGHT_DECAY = 1e-5  # L2 regularization
+GRADIENT_CLIP = 1.0  # Gradient clipping
 
 # --- 1. Load and Parse Data ---
 def load_and_parse_data_from_string(training_data_text: str):
@@ -87,24 +92,132 @@ class TalmudDataset(Dataset):
 
 # --- 4. Model Definition ---
 class TalmudClassifierLSTM(nn.Module):
-    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim):
+    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, num_layers=2):
         super(TalmudClassifierLSTM, self).__init__()
         self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
-        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, dropout=0.3, num_layers=2)
-        self.dropout = nn.Dropout(0.5)
-        self.fc1 = nn.Linear(hidden_dim, 64)
+        # Bidirectional LSTM - uses both forward and backward contexts
+        self.lstm = nn.LSTM(
+            embedding_dim, 
+            hidden_dim // 2,  # Divide by 2 because bidirectional doubles the output
+            batch_first=True, 
+            dropout=0.3 if num_layers > 1 else 0,
+            num_layers=num_layers,
+            bidirectional=True
+        )
+        self.dropout1 = nn.Dropout(0.5)
+        self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2)
         self.relu = nn.ReLU()
-        self.fc2 = nn.Linear(64, output_dim)
+        self.dropout2 = nn.Dropout(0.3)
+        self.fc2 = nn.Linear(hidden_dim // 2, output_dim)
 
     def forward(self, text):
         embedded = self.embedding(text)
-        _, (hidden, _) = self.lstm(embedded)
-        hidden = self.dropout(hidden[-1])
+        # Get LSTM output - use both forward and backward hidden states
+        lstm_out, (hidden, _) = self.lstm(embedded)
+        # Concatenate forward and backward hidden states from last layer
+        # hidden shape: (num_layers * num_directions, batch, hidden_size)
+        if self.lstm.bidirectional:
+            hidden_forward = hidden[-2]
+            hidden_backward = hidden[-1]
+            hidden = torch.cat([hidden_forward, hidden_backward], dim=1)
+        else:
+            hidden = hidden[-1]
+        
+        hidden = self.dropout1(hidden)
         out = self.fc1(hidden)
         out = self.relu(out)
+        out = self.dropout2(out)
         out = self.fc2(out)
         return out
 
+# --- 4.5. Helper Functions ---
+def calculate_class_weights(labels, label_encoder):
+    """Calculate class weights for weighted loss function."""
+    # Count occurrences of each class
+    label_counts = Counter(labels)
+    total_samples = len(labels)
+    num_classes = len(label_encoder.classes_)
+    
+    # Calculate weights: inverse frequency, normalized
+    weights = np.ones(num_classes)
+    for i, class_name in enumerate(label_encoder.classes_):
+        count = label_counts.get(class_name, 1)  # Avoid division by zero
+        # Weight is inversely proportional to frequency
+        weights[i] = total_samples / (num_classes * count)
+    
+    # Normalize weights to sum to num_classes
+    weights = weights / weights.sum() * num_classes
+    return torch.FloatTensor(weights)
+
+def create_weighted_sampler(labels, label_encoder):
+    """Create a weighted sampler for balanced batch sampling."""
+    # Convert string labels to encoded labels
+    encoded_labels = label_encoder.transform(labels)
+    
+    # Calculate weights for each sample
+    label_counts = Counter(encoded_labels)
+    total_samples = len(encoded_labels)
+    num_classes = len(label_encoder.classes_)
+    
+    sample_weights = np.ones(total_samples)
+    for i, label in enumerate(encoded_labels):
+        count = label_counts[label]
+        # Weight inversely proportional to class frequency
+        sample_weights[i] = total_samples / (num_classes * count)
+    
+    return WeightedRandomSampler(
+        weights=sample_weights,
+        num_samples=len(sample_weights),
+        replacement=True
+    )
+
+def evaluate_model(model, data_loader, criterion, label_encoder, device='cpu'):
+    """Evaluate model and return metrics."""
+    model.eval()
+    all_predicted = []
+    all_labels = []
+    total_loss = 0.0
+    num_batches = 0
+    
+    with torch.no_grad():
+        for sequences, labels in data_loader:
+            sequences = sequences.to(device)
+            labels = labels.to(device)
+            
+            outputs = model(sequences)
+            loss = criterion(outputs, labels)
+            
+            total_loss += loss.item()
+            num_batches += 1
+            
+            _, predicted = torch.max(outputs.data, 1)
+            all_predicted.extend(predicted.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+    
+    avg_loss = total_loss / num_batches if num_batches > 0 else 0.0
+    accuracy = 100 * np.mean(np.array(all_predicted) == np.array(all_labels))
+    
+    # Calculate per-class F1 scores
+    label_names = label_encoder.classes_
+    f1_scores_dict = {}
+    for i, label_name in enumerate(label_names):
+        binary_true = np.array(all_labels) == i
+        binary_pred = np.array(all_predicted) == i
+        f1 = f1_score(binary_true, binary_pred, zero_division=0)
+        f1_scores_dict[label_name] = float(f1)
+    
+    # Calculate macro-averaged F1 score
+    macro_f1 = np.mean(list(f1_scores_dict.values()))
+    
+    return {
+        'accuracy': accuracy,
+        'loss': avg_loss,
+        'f1_scores': f1_scores_dict,
+        'macro_f1': macro_f1,
+        'predictions': all_predicted,
+        'labels': all_labels
+    }
+
 # --- 5. Training Function ---
 def train_model(training_data_text: str):
     """
@@ -148,17 +261,31 @@ def train_model(training_data_text: str):
     print(f"\nTotal samples: {len(all_texts)}")
     print(f"Training set size: {len(train_texts)} (80%)")
     print(f"Test set size: {len(test_texts)} (20%)")
+    
+    # Print class distribution
+    train_label_counts = Counter(train_labels)
+    print("\nTraining set class distribution:")
+    for label, count in sorted(train_label_counts.items()):
+        print(f"  {label}: {count} ({100*count/len(train_labels):.1f}%)")
 
     # Build vocabulary and label encoder ONLY on the training data
     word_to_idx = build_vocab(train_texts, VOCAB_SIZE)
     label_encoder = LabelEncoder()
     label_encoder.fit(train_labels)
     num_classes = len(label_encoder.classes_)
+    
+    # Calculate class weights for weighted loss
+    class_weights = calculate_class_weights(train_labels, label_encoder)
+    print(f"\nClass weights: {dict(zip(label_encoder.classes_, class_weights.numpy()))}")
+    
+    # Set device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
 
     # Set up K-Fold Cross-Validation
     kfold = KFold(n_splits=N_SPLITS, shuffle=True, random_state=42)
     
-    best_val_accuracy = 0.0
+    best_val_macro_f1 = 0.0
     best_model_state = None
     fold_results = []
 
@@ -171,51 +298,118 @@ def train_model(training_data_text: str):
         print(f"\n----- FOLD {fold+1}/{N_SPLITS} -----")
 
         # Create data subsets for the current fold
-        train_subsampler = torch.utils.data.SubsetRandomSampler(train_ids)
-        val_subsampler = torch.utils.data.SubsetRandomSampler(val_ids)
-
-        train_loader = DataLoader(full_train_dataset, batch_size=BATCH_SIZE, sampler=train_subsampler)
-        val_loader = DataLoader(full_train_dataset, batch_size=BATCH_SIZE, sampler=val_subsampler)
+        train_subset_texts = [train_texts[i] for i in train_ids]
+        train_subset_labels = [train_labels[i] for i in train_ids]
+        val_subset_texts = [train_texts[i] for i in val_ids]
+        val_subset_labels = [train_labels[i] for i in val_ids]
+        
+        # Create datasets for this fold
+        train_dataset = TalmudDataset(train_subset_texts, train_subset_labels, word_to_idx, label_encoder, MAX_LEN)
+        val_dataset = TalmudDataset(val_subset_texts, val_subset_labels, word_to_idx, label_encoder, MAX_LEN)
+        
+        # Create weighted sampler for balanced training
+        weighted_sampler = create_weighted_sampler(train_subset_labels, label_encoder)
+        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, sampler=weighted_sampler)
+        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)
         
         # Initialize a new model for each fold
         model = TalmudClassifierLSTM(len(word_to_idx), EMBEDDING_DIM, HIDDEN_DIM, num_classes)
-        criterion = nn.CrossEntropyLoss()
-        optimizer = optim.Adam(model.parameters(), lr=0.001)
+        model = model.to(device)
+        
+        # Use weighted loss to handle class imbalance
+        class_weights_device = class_weights.to(device)
+        criterion = nn.CrossEntropyLoss(weight=class_weights_device)
+        
+        # Optimizer with weight decay for regularization
+        optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
+        
+        # Learning rate scheduler - reduce LR on plateau
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode='max', factor=0.5, patience=3
+        )
+        
+        # Early stopping variables
+        best_fold_macro_f1 = 0.0
+        best_fold_model_state = None
+        patience_counter = 0
 
+        # Training loop with early stopping
         for epoch in range(NUM_EPOCHS):
             model.train()
+            epoch_loss = 0.0
+            num_batches = 0
+            
             for sequences, labels in train_loader:
+                sequences = sequences.to(device)
+                labels = labels.to(device)
+                
                 optimizer.zero_grad()
                 outputs = model(sequences)
                 loss = criterion(outputs, labels)
                 loss.backward()
+                
+                # Gradient clipping to prevent exploding gradients
+                torch.nn.utils.clip_grad_norm_(model.parameters(), GRADIENT_CLIP)
+                
                 optimizer.step()
-
-        # Evaluate the model on the validation set for this fold
-        model.eval()
-        all_predicted = []
-        all_labels_val = []
+                epoch_loss += loss.item()
+                num_batches += 1
+            
+            avg_epoch_loss = epoch_loss / num_batches if num_batches > 0 else 0.0
+            
+            # Evaluate on validation set
+            val_metrics = evaluate_model(model, val_loader, criterion, label_encoder, device)
+            
+            # Update learning rate based on validation macro F1
+            scheduler.step(val_metrics['macro_f1'])
+            
+            # Print progress
+            print(f"Epoch {epoch+1}/{NUM_EPOCHS} - Loss: {avg_epoch_loss:.4f}, "
+                  f"Val Acc: {val_metrics['accuracy']:.2f}%, "
+                  f"Val Macro F1: {val_metrics['macro_f1']:.4f}")
+            print(f"  Per-class F1: {', '.join([f'{k}: {v:.3f}' for k, v in val_metrics['f1_scores'].items()])}")
+            
+            # Early stopping based on macro F1 score
+            if val_metrics['macro_f1'] > best_fold_macro_f1:
+                best_fold_macro_f1 = val_metrics['macro_f1']
+                best_fold_model_state = copy.deepcopy(model.state_dict())
+                patience_counter = 0
+            else:
+                patience_counter += 1
+                if patience_counter >= EARLY_STOPPING_PATIENCE:
+                    print(f"Early stopping triggered at epoch {epoch+1}")
+                    break
         
-        with torch.no_grad():
-            for sequences, labels in val_loader:
-                outputs = model(sequences)
-                _, predicted = torch.max(outputs.data, 1)
-                all_predicted.extend(predicted.cpu().numpy())
-                all_labels_val.extend(labels.cpu().numpy())
-
-        accuracy = 100 * np.mean(np.array(all_predicted) == np.array(all_labels_val))
-        fold_results.append(accuracy)
-        print(f"Validation Accuracy for Fold {fold+1}: {accuracy:.2f}%")
+        # Load best model for this fold
+        if best_fold_model_state is not None:
+            model.load_state_dict(best_fold_model_state)
+        
+        # Final evaluation on validation set
+        val_metrics = evaluate_model(model, val_loader, criterion, label_encoder, device)
+        fold_results.append({
+            'accuracy': val_metrics['accuracy'],
+            'macro_f1': val_metrics['macro_f1'],
+            'f1_scores': val_metrics['f1_scores']
+        })
+        
+        print(f"\nFold {fold+1} Results:")
+        print(f"  Validation Accuracy: {val_metrics['accuracy']:.2f}%")
+        print(f"  Validation Macro F1: {val_metrics['macro_f1']:.4f}")
+        for label, f1 in val_metrics['f1_scores'].items():
+            print(f"  {label} F1: {f1:.4f}")
 
-        # Save the best model found across all folds
-        # Always save at least the first fold's model, or if this fold is better
-        if best_model_state is None or accuracy >= best_val_accuracy:
-            best_val_accuracy = accuracy
+        # Save the best model found across all folds (based on macro F1)
+        if best_model_state is None or val_metrics['macro_f1'] >= best_val_macro_f1:
+            best_val_macro_f1 = val_metrics['macro_f1']
             best_model_state = copy.deepcopy(model.state_dict())
 
     print("\n----- Cross-Validation Summary -----")
-    print(f"Fold Accuracies: {[f'{acc:.2f}%' for acc in fold_results]}")
-    print(f"Average CV Accuracy: {np.mean(fold_results):.2f}%")
+    acc_strs = [f"{r['accuracy']:.2f}%" for r in fold_results]
+    f1_strs = [f"{r['macro_f1']:.4f}" for r in fold_results]
+    print(f"Fold Accuracies: {acc_strs}")
+    print(f"Fold Macro F1s: {f1_strs}")
+    print(f"Average CV Accuracy: {np.mean([r['accuracy'] for r in fold_results]):.2f}%")
+    print(f"Average CV Macro F1: {np.mean([r['macro_f1'] for r in fold_results]):.4f}")
 
     # Verify that we have a model state to load
     if best_model_state is None:
@@ -225,45 +419,38 @@ def train_model(training_data_text: str):
     print("\n----- Final Evaluation on Test Set -----")
     final_model = TalmudClassifierLSTM(len(word_to_idx), EMBEDDING_DIM, HIDDEN_DIM, num_classes)
     final_model.load_state_dict(best_model_state)
-    final_model.eval()
+    final_model = final_model.to(device)
 
     test_dataset = TalmudDataset(test_texts, test_labels, word_to_idx, label_encoder, MAX_LEN)
-    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE)
-
-    all_test_predicted = []
-    all_test_labels = []
-    test_losses = []
+    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
     
-    criterion = nn.CrossEntropyLoss()
+    # Use weighted loss for evaluation too
+    class_weights_device = class_weights.to(device)
+    criterion = nn.CrossEntropyLoss(weight=class_weights_device)
     
-    with torch.no_grad():
-        for sequences, labels in test_loader:
-            outputs = final_model(sequences)
-            loss = criterion(outputs, labels)
-            test_losses.append(loss.item())
-            _, predicted = torch.max(outputs.data, 1)
-            all_test_predicted.extend(predicted.cpu().numpy())
-            all_test_labels.extend(labels.cpu().numpy())
+    # Evaluate on test set
+    test_metrics = evaluate_model(final_model, test_loader, criterion, label_encoder, device)
     
-    test_accuracy = 100 * np.mean(np.array(all_test_predicted) == np.array(all_test_labels))
-    avg_loss = np.mean(test_losses)
+    test_accuracy = test_metrics['accuracy']
+    avg_loss = test_metrics['loss']
+    f1_scores_dict = test_metrics['f1_scores']
+    macro_f1 = test_metrics['macro_f1']
     
     print(f"Accuracy on the unseen test set: {test_accuracy:.2f}%")
     print(f"Average loss: {avg_loss:.4f}")
+    print(f"Macro-averaged F1 score: {macro_f1:.4f}")
+    print("\nPer-class F1 scores:")
+    for label_name, f1 in f1_scores_dict.items():
+        print(f"  {label_name}: {f1:.4f}")
     
-    # Calculate F1 scores per category
-    label_names = label_encoder.classes_
-    f1_scores_dict = {}
-    
-    for i, label_name in enumerate(label_names):
-        # Create binary labels for this category
-        binary_true = np.array(all_test_labels) == i
-        binary_pred = np.array(all_test_predicted) == i
-        
-        # Calculate F1 score
-        f1 = f1_score(binary_true, binary_pred, zero_division=0)
-        f1_scores_dict[label_name] = float(f1)
-        print(f"F1 Score for {label_name}: {f1:.4f}")
+    # Print detailed classification report
+    print("\nClassification Report:")
+    print(classification_report(
+        test_metrics['labels'],
+        test_metrics['predictions'],
+        target_names=label_encoder.classes_,
+        zero_division=0
+    ))
     
     # Convert accuracy to 0-1 range for callback
     accuracy_normalized = test_accuracy / 100.0
@@ -274,8 +461,11 @@ def train_model(training_data_text: str):
         word_to_idx_path = '/tmp/word_to_idx.pt'
         label_encoder_path = '/tmp/label_encoder.pkl'
         
+        # Move model to CPU for saving (to ensure compatibility)
+        final_model_cpu = final_model.cpu()
+        
         # Save model state dict
-        torch.save(final_model.state_dict(), model_path)
+        torch.save(final_model_cpu.state_dict(), model_path)
         print(f"Saved model to {model_path}")
         
         # Save word_to_idx dictionary
@@ -287,6 +477,9 @@ def train_model(training_data_text: str):
             pickle.dump(label_encoder, f)
         print(f"Saved label_encoder to {label_encoder_path}")
         
+        # Move model back to device for return
+        final_model = final_model.to(device)
+        
     except Exception as e:
         print(f"Warning: Failed to save model artifacts to /tmp: {e}")
         # Continue even if saving fails - model is still returned in result
@@ -300,6 +493,7 @@ def train_model(training_data_text: str):
             'accuracy': accuracy_normalized,
             'loss': float(avg_loss),
             'f1_scores': f1_scores_dict,
+            'macro_f1': float(macro_f1),
             'model_path': '/tmp/latest_model.pt'  # Path to saved model
         }
     }