feat: Implement Learnable Gated Pooling model

requirements.txt

@@ -1 +1,3 @@
-numpy
+numpy>=1.21.0
+torch>=2.0.0
+scikit-learn>=1.0.0

src/evaluate.py

@@ -0,0 +1,36 @@
+import torch
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+
+def evaluate_model(model, test_loader, device):
+    """
+    Evaluate the trained model on test data
+    """
+    model.eval()
+    all_preds = []
+    all_targets = []
+
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            pred = output.argmax(dim=1).cpu().numpy()
+            all_preds.extend(pred)
+            all_targets.extend(target.cpu().numpy())
+
+    # Calculate metrics
+    accuracy = accuracy_score(all_targets, all_preds)
+    precision, recall, f1, _ = precision_recall_fscore_support(
+        all_targets,
+        all_preds,
+        average='weighted',
+        zero_division=0  # Handle undefined recall
+    )
+
+    metrics = {
+        'accuracy': accuracy,
+        'precision': precision,
+        'recall': recall,
+        'f1': f1
+    }
+
+    return metrics

src/main.py

@@ -1 +1,66 @@
-print('Hello, World!')
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import argparse
+from models import LearnableGatedPooling
+from train import train_model
+from evaluate import evaluate_model
+from preprocess import prepare_data
+
+def main(args):
+    # Set device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Prepare data
+    train_loader, val_loader, test_loader = prepare_data(
+        args.data_path,
+        batch_size=args.batch_size
+    )
+
+    # Initialize model
+    model = LearnableGatedPooling(
+        input_dim=args.input_dim,
+        seq_len=args.seq_len,
+        num_classes=2  # Binary classification
+    )
+
+    # Define loss function and optimizer
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+
+    # Train model
+    model, history = train_model(
+        model=model,
+        train_loader=train_loader,
+        val_loader=val_loader,
+        criterion=criterion,
+        optimizer=optimizer,
+        num_epochs=args.num_epochs,
+        device=device
+    )
+
+    # Evaluate model
+    metrics = evaluate_model(model, test_loader, device)
+
+    # Print results
+    print("\nEvaluation Results:")
+    for metric_name, value in metrics.items():
+        print(f"{metric_name}: {value:.4f}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Train and evaluate LearnableGatedPooling model')
+    parser.add_argument('--data_path', type=str, default='data',
+                        help='path to data directory')
+    parser.add_argument('--input_dim', type=int, default=768,
+                        help='input dimension size')
+    parser.add_argument('--seq_len', type=int, default=10,
+                        help='sequence length')
+    parser.add_argument('--batch_size', type=int, default=32,
+                        help='batch size for training')
+    parser.add_argument('--learning_rate', type=float, default=0.001,
+                        help='learning rate')
+    parser.add_argument('--num_epochs', type=int, default=10,
+                        help='number of training epochs')
+
+    args = parser.parse_args()
+    main(args)

src/models.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LearnableGatedPooling(nn.Module):
+    def __init__(self, input_dim, seq_len, num_classes=2):
+        super(LearnableGatedPooling, self).__init__()
+        self.weights = nn.Parameter(torch.ones(input_dim))
+        self.gate_linear = nn.Linear(input_dim, 1)  # Linear layer for gating
+        self.classifier = nn.Linear(input_dim, num_classes)  # Classification head
+
+    def forward(self, x):
+        # x: (batch_size, seq_len, input_dim)
+        weighted_x = x * self.weights
+        gate_values = torch.sigmoid(self.gate_linear(x)).squeeze(2)  # (batch_size, seq_len)
+        gated_x = weighted_x * gate_values.unsqueeze(2)
+        pooled_vector = torch.mean(gated_x, dim=1)  # Average pooling
+        logits = self.classifier(pooled_vector)  # (batch_size, num_classes)
+        return logits

src/preprocess.py

@@ -0,0 +1,53 @@
+import torch
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+
+class SequenceDataset(Dataset):
+    """
+    Dataset class for handling sequence data
+    """
+    def __init__(self, sequences, labels):
+        self.sequences = torch.FloatTensor(sequences)
+        self.labels = torch.LongTensor(labels)
+
+    def __len__(self):
+        return len(self.sequences)
+
+    def __getitem__(self, idx):
+        return self.sequences[idx], self.labels[idx]
+
+def prepare_data(data_path, batch_size=32):
+    """
+    Prepare data loaders for training, validation, and testing
+    """
+    # Load and preprocess data
+    # Note: Implement actual data loading logic based on your data format
+
+    # Example data creation (replace with actual data loading)
+    def create_dummy_data(num_samples, seq_len, input_dim):
+        half_samples = num_samples // 2
+        sequences_0 = np.random.randn(half_samples, seq_len, input_dim)
+        labels_0 = np.zeros(half_samples, dtype=int)
+        sequences_1 = np.random.randn(num_samples - half_samples, seq_len, input_dim)
+        labels_1 = np.ones(num_samples - half_samples, dtype=int)
+
+        sequences = np.concatenate([sequences_0, sequences_1], axis=0)
+        labels = np.concatenate([labels_0, labels_1], axis=0)
+        return sequences, labels
+
+    # Create datasets
+    train_sequences, train_labels = create_dummy_data(1000, 10, 768)
+    val_sequences, val_labels = create_dummy_data(200, 10, 768)
+    test_sequences, test_labels = create_dummy_data(200, 10, 768)
+
+    # Create dataset objects
+    train_dataset = SequenceDataset(train_sequences, train_labels)
+    val_dataset = SequenceDataset(val_sequences, val_labels)
+    test_dataset = SequenceDataset(test_sequences, test_labels)
+
+    # Create data loaders
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size)
+
+    return train_loader, val_loader, test_loader

src/train.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from models import LearnableGatedPooling
+
+def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device):
+    """
+    Train the LearnableGatedPooling model
+    """
+    model = model.to(device)
+    history = {'train_loss': [], 'val_loss': []}
+
+    for epoch in range(num_epochs):
+        model.train()
+        train_loss = 0.0
+        for batch_idx, (data, target) in enumerate(train_loader):
+            data, target = data.to(device), target.to(device)
+
+            optimizer.zero_grad()
+            output = model(data)
+            loss = criterion(output, target)
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+
+        # Validation phase
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for data, target in val_loader:
+                data, target = data.to(device), target.to(device)
+                output = model(data)
+                val_loss += criterion(output, target).item()
+
+        # Record losses
+        train_loss /= len(train_loader)
+        val_loss /= len(val_loader)
+        history['train_loss'].append(train_loss)
+        history['val_loss'].append(val_loss)
+
+        print(f'Epoch {epoch+1}/{num_epochs}:')
+        print(f'Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')
+
+    return model, history