diff --git a/requirements.txt b/requirements.txt index 24ce15a..cef06d7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ numpy +torch +torchvision diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..d8ee75e --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +# Make src directory a Python package diff --git a/src/evaluate.py b/src/evaluate.py index e69de29..9eec6be 100644 --- a/src/evaluate.py +++ b/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() diff --git a/src/main.py b/src/main.py index 0990720..52b7b64 100644 --- a/src/main.py +++ b/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() diff --git a/src/model_components.py b/src/model_components.py new file mode 100644 index 0000000..59aa5f0 --- /dev/null +++ b/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 diff --git a/src/preprocess.py b/src/preprocess.py index e69de29..c922e54 100644 --- a/src/preprocess.py +++ b/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 diff --git a/src/train.py b/src/train.py index e69de29..ba1d655 100644 --- a/src/train.py +++ b/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