feat: Implement Learnable Gated Pooling approach

requirements.txt

@@ -1 +1,3 @@
 numpy
+torch
+torchvision

src/__init__.py

@@ -0,0 +1 @@
+# Make src directory a Python package

src/evaluate.py

@@ -0,0 +1,22 @@
+import torch
+from src.preprocess import load_and_preprocess_data
+from src.model_components import LearnableGatedPooling
+
+def evaluate_model(pooling_layer, classifier):
+    """
+    Evaluates the trained model using accuracy metric.
+
+    Args:
+        pooling_layer: Trained LearnableGatedPooling instance
+        classifier: Trained classifier layer
+
+    Returns:
+        float: Accuracy of the model on the evaluation data
+    """
+    embeddings, labels = load_and_preprocess_data()
+    with torch.no_grad():
+        pooled_output = pooling_layer(embeddings)
+        logits = classifier(pooled_output)
+        predictions = torch.argmax(logits, dim=1)
+        accuracy = (predictions == labels).float().mean()
+    return accuracy.item()

src/main.py

@@ -1 +1,20 @@
-print('Hello, World!')
+from src.train import train_model
+from src.evaluate import evaluate_model
+
+def run_experiment():
+    """
+    Runs the complete experiment workflow:
+    1. Trains the model using LearnableGatedPooling
+    2. Evaluates the model and reports accuracy
+    """
+    # Train
+    print("Starting training...")
+    trained_pooling, trained_classifier = train_model()
+
+    # Evaluate
+    print("\nStarting evaluation...")
+    accuracy = evaluate_model(trained_pooling, trained_classifier)
+    print(f"Evaluation Accuracy: {accuracy:.4f}")
+
+if __name__ == "__main__":
+    run_experiment()

src/model_components.py

@@ -0,0 +1,17 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LearnableGatedPooling(nn.Module):
+    def __init__(self, input_dim):
+        super(LearnableGatedPooling, self).__init__()
+        self.weights = nn.Parameter(torch.ones(input_dim))
+        self.gate_linear = nn.Linear(input_dim, 1)
+
+    def forward(self, x):
+        # x shape: (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 operation
+        return pooled_vector

src/preprocess.py

@@ -0,0 +1,21 @@
+import torch
+
+def load_and_preprocess_data(batch_size=32, seq_len=10, input_dim=768):
+    """
+    Simulates loading and preprocessing of data for the LearnableGatedPooling model.
+
+    Args:
+        batch_size (int): Number of samples in a batch
+        seq_len (int): Length of the input sequence
+        input_dim (int): Dimension of each vector in the sequence
+
+    Returns:
+        tuple: (embeddings, labels)
+            - embeddings: torch.Tensor of shape (batch_size, seq_len, input_dim)
+            - labels: torch.Tensor of shape (batch_size,) with binary values
+    """
+    # Placeholder that simulates loading data
+    # In a real scenario, load data from `data/` directory and preprocess as needed
+    embeddings = torch.randn(batch_size, seq_len, input_dim)
+    labels = torch.randint(0, 2, (batch_size,))
+    return embeddings, labels

src/train.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from src.preprocess import load_and_preprocess_data
+from src.model_components import LearnableGatedPooling
+
+def train_model():
+    """
+    Trains the LearnableGatedPooling model with a simple classifier.
+    Returns the trained model components.
+    """
+
+    # Hyperparameters
+    input_dim = 768
+    learning_rate = 1e-3
+    num_epochs = 2
+
+    # Prepare data
+    embeddings, labels = load_and_preprocess_data()
+
+    # Define model
+    pooling_layer = LearnableGatedPooling(input_dim)
+    classifier = nn.Linear(input_dim, 2)
+
+    # Optimizer
+    optimizer = optim.Adam(list(pooling_layer.parameters()) + list(classifier.parameters()), lr=learning_rate)
+    loss_fn = nn.CrossEntropyLoss()
+
+    # Training loop
+    for epoch in range(num_epochs):
+        optimizer.zero_grad()
+        # Forward pass
+        pooled_output = pooling_layer(embeddings)
+        logits = classifier(pooled_output)
+        loss = loss_fn(logits, labels)
+
+        # Backward pass and optimization
+        loss.backward()
+        optimizer.step()
+
+        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
+
+    return pooling_layer, classifier