neural_net.py

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.metrics import f1_score
import pandas as pd
import numpy as np
import ast
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from tqdm import tqdm
import matplotlib.pyplot as plt
import itertools
from sklearn.metrics import f1_score

df = pd.read_csv("data/questions_tags_embedded.csv")
df['target'] = df['target'].apply(ast.literal_eval)

X = np.array(df['embedded_question'].apply(lambda x: np.fromstring(x.strip('[]'), sep=' ')).tolist())

mlb = MultiLabelBinarizer()
# Use df['languagues'] for programming language classification
# Y = mlb.fit_transform(df['languages'])

# Use df['target'] for multi label tag classification
Y = mlb.fit_transform(df["target"])

X_train, X_temp, Y_train, Y_temp = train_test_split(X, Y, test_size=0.2, random_state=42)
X_val, X_test, Y_val, Y_test = train_test_split(X_temp, Y_temp, test_size=0.5, random_state=42)

print("Train:", X_train.shape, " Val:", X_val.shape, " Test:", X_test.shape)


class EmbeddingDataset(Dataset):
    def __init__(self, X, Y):
        self.X = torch.tensor(X, dtype=torch.float32)
        self.Y = torch.tensor(Y, dtype=torch.float32)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.Y[idx]

train_loader = DataLoader(EmbeddingDataset(X_train, Y_train), batch_size=128, shuffle=True)
val_loader   = DataLoader(EmbeddingDataset(X_val, Y_val), batch_size=128)
test_loader  = DataLoader(EmbeddingDataset(X_test, Y_test), batch_size=128)


# Model
class MultiLabelClassifier(nn.Module):
    def __init__(self, input_dim, output_dim, hidden1=512, hidden2=256, dropout=0.3):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, hidden1),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden1, hidden2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden2, output_dim),
        )

    def forward(self, x):
        return self.model(x)
    
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
criterion = nn.BCEWithLogitsLoss()

def train_model(model, train_loader, val_loader, optimizer, epochs=20, threshold=0.5):
    train_losses = []
    val_losses = []
    val_f1_scores = []

    for epoch in range(epochs):
        print(f"\nEpoch {epoch+1}/{epochs}")

        # training
        model.train()
        total_train_loss = 0

        for batch_x, batch_y in tqdm(train_loader):
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)

            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_train_loss += loss.item()

        model.eval()
        total_val_loss = 0
        preds, labels = [], []

        with torch.no_grad():
            for batch_x, batch_y in val_loader:
                batch_x, batch_y = batch_x.to(device), batch_y.to(device)
                outputs = model(batch_x)

                loss = criterion(outputs, batch_y)
                total_val_loss += loss.item()

                out = torch.sigmoid(outputs).cpu().numpy()
                preds.append((out > threshold).astype(int))
                labels.append(batch_y.cpu().numpy())

        preds = np.vstack(preds)
        labels = np.vstack(labels)
        f1 = f1_score(labels, preds, average='micro')

        train_losses.append(total_train_loss)
        val_losses.append(total_val_loss)
        val_f1_scores.append(f1)

        print(f"  Train Loss: {total_train_loss:.4f}")
        print(f"  Val Loss:   {total_val_loss:.4f}")
        print(f"  Val F1:     {f1:.4f}")

    return train_losses, val_losses, val_f1_scores


# model evaluation
def evaluate(model, test_loader, threshold=0.5):
    model.eval()
    preds, labels = [], []

    with torch.no_grad():
        for batch_x, batch_y in tqdm(test_loader):
            batch_x = batch_x.to(device)
            outputs = torch.sigmoid(model(batch_x)).cpu().numpy()

            preds.append((outputs > threshold).astype(int))
            labels.append(batch_y.numpy())

    preds = np.vstack(preds)
    labels = np.vstack(labels)

    f1 = f1_score(labels, preds, average='micro')

    return f1

input_dim = X_train.shape[1]
output_dim = Y_train.shape[1]

model = MultiLabelClassifier(
    input_dim=input_dim,
    output_dim=output_dim,
    hidden1=512,
    hidden2=256,
    dropout=0.2
).to(device)
    
optimizer = optim.Adam(model.parameters(), lr=1e-4)
train_losses, val_losses, val_f1_scores = train_model(
    model=model,
    train_loader=train_loader,
    val_loader=val_loader,
    optimizer=optimizer,
    epochs=50
)

print(evaluate(model=model, test_loader=test_loader))

# Plot curves
plt.figure(figsize=(12, 8))
plt.plot(train_losses, label="Train Loss")
plt.plot(val_losses, label="Validation Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.title("Training vs Validation Loss")
plt.legend()
plt.grid(True)
plt.savefig('Training Curves.pdf')
plt.show()

plt.figure(figsize=(12,8))
plt.plot(val_f1_scores)
plt.title("Validation F1 Score per Epoch")
plt.xlabel("Epoch")
plt.ylabel("F1 Score (micro)")
plt.grid(True)
plt.savefig('F1 Validation Scores.pdf')
plt.show()