The solution was based on this notebook (training an LSTM), along with additional hyperparameter tuning. Another approach worth exploring would be to use a transformer instead.
import os
import math
import random
from collections import Counter
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")
def tokenize(program: str):
program = str(program).strip()
program = program.replace(";", " ; ")
return program.split()
train_df["tokens"] = train_df["program"].apply(tokenize)
test_df["tokens"] = test_df["program"].apply(tokenize)
counter = Counter()
for toks in pd.concat([train_df["tokens"], test_df["tokens"]], axis=0):
counter.update(toks)
PAD = ""
UNK = ""
stoi = {PAD: 0, UNK: 1}
for tok, _ in counter.most_common():
if tok not in stoi:
stoi[tok] = len(stoi)
itos = {i: t for t, i in stoi.items()}
def encode(tokens):
return [stoi.get(tok, stoi[UNK]) for tok in tokens]
train_df["ids"] = train_df["tokens"].apply(encode)
test_df["ids"] = test_df["tokens"].apply(encode)
train_idx, val_idx = train_test_split(
np.arange(len(train_df)),
test_size=0.1,
random_state=SEED
)
train_part = train_df.iloc[train_idx].reset_index(drop=True)
val_part = train_df.iloc[val_idx].reset_index(drop=True)
y_mean = train_part["output"].mean()
y_std = train_part["output"].std()
if y_std == 0:
y_std = 1.0
def normalize_y(y):
return (y - y_mean) / y_std
def denormalize_y(y):
return y * y_std + y_mean
class ProgramDataset(Dataset):
def __init__(self, df, has_target=True):
self.seqs = df["ids"].tolist()
self.has_target = has_target
self.targets = df["output"].tolist() if has_target else None
def __len__(self):
return len(self.seqs)
def __getitem__(self, idx):
x = self.seqs[idx]
if self.has_target:
return x, self.targets[idx]
return x
def collate_train(batch):
seqs, ys = zip(*batch)
lengths = torch.tensor([len(s) for s in seqs], dtype=torch.long)
max_len = max(lengths).item()
x = torch.zeros(len(seqs), max_len, dtype=torch.long)
for i, s in enumerate(seqs):
x[i, :len(s)] = torch.tensor(s, dtype=torch.long)
y = torch.tensor([normalize_y(float(v)) for v in ys], dtype=torch.float32)
return x, lengths, y
def collate_test(batch):
lengths = torch.tensor([len(s) for s in batch], dtype=torch.long)
max_len = max(lengths).item()
x = torch.zeros(len(batch), max_len, dtype=torch.long)
for i, s in enumerate(batch):
x[i, :len(s)] = torch.tensor(s, dtype=torch.long)
return x, lengths
train_loader = DataLoader(
ProgramDataset(train_part, has_target=True),
batch_size=64,
shuffle=True,
collate_fn=collate_train,
num_workers=0
)
val_loader = DataLoader(
ProgramDataset(val_part, has_target=True),
batch_size=128,
shuffle=False,
collate_fn=collate_train,
num_workers=0
)
test_loader = DataLoader(
ProgramDataset(test_df, has_target=False),
batch_size=128,
shuffle=False,
collate_fn=collate_test,
num_workers=0
)
class ProgramRegressor(nn.Module):
def __init__(
self,
vocab_size,
emb_dim=64,
hidden_dim=64,
num_layers=2,
dropout=0.01
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
self.gru = nn.GRU(
emb_dim,
hidden_dim,
num_layers=num_layers,
batch_first=True,
bidirectional=True,
dropout=dropout if num_layers > 1 else 0.0,
)
feat_dim = hidden_dim * 4
self.head = nn.Sequential(
nn.Linear(feat_dim, 512),
nn.LayerNorm(512),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(512, 512),
nn.LayerNorm(512),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(512, 256),
nn.LayerNorm(256),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(128, 64),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(64, 1),
)
def forward(self, x, lengths):
emb = self.embedding(x)
packed = nn.utils.rnn.pack_padded_sequence(
emb, lengths.cpu(), batch_first=True, enforce_sorted=False
)
packed_out, _ = self.gru(packed)
out, _ = nn.utils.rnn.pad_packed_sequence(packed_out, batch_first=True)
mask = (x != 0).unsqueeze(-1)
out = out[:, :x.size(1), :]
out_masked = out * mask
denom = mask.sum(dim=1).clamp(min=1)
mean_pool = out_masked.sum(dim=1) / denom
max_pool = out.masked_fill(~mask, -1e9).max(dim=1).values
feat = torch.cat([mean_pool, max_pool], dim=1)
return self.head(feat).squeeze(-1)
model = ProgramRegressor(vocab_size=len(stoi)).to(device)
criterion = nn.SmoothL1Loss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# Consider also adding a scheduler like:
"""scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode="min", factor=0.5, patience=2
)"""
def evaluate(loader):
model.eval()
preds, trues = [], []
with torch.no_grad():
for x, lengths, y in loader:
x = x.to(device)
lengths = lengths.to(device)
y = y.to(device)
pred = model(x, lengths)
preds.append(pred.cpu().numpy())
trues.append(y.cpu().numpy())
preds = np.concatenate(preds)
trues = np.concatenate(trues)
preds_real = denormalize_y(preds)
trues_real = denormalize_y(trues)
return mean_absolute_error(trues_real, preds_real)
def fit_one_epoch(loader):
model.train()
total_loss = 0.0
n = 0
for x, lengths, y in loader:
x = x.to(device)
lengths = lengths.to(device)
y = y.to(device)
optimizer.zero_grad()
pred = model(x, lengths)
loss = criterion(pred, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
bs = x.size(0)
total_loss += loss.item() * bs
n += bs
return total_loss / n
best_val = float("inf")
best_state = None
patience = 20
bad_epochs = 0
num_epochs = 50
for epoch in range(1, num_epochs + 1):
train_loss = fit_one_epoch(train_loader)
val_mae = evaluate(val_loader)
print(f"Epoch {epoch:02d} | train loss: {train_loss:.5f} | val MAE: {val_mae:.5f}")
if val_mae < best_val:
best_val = val_mae
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
bad_epochs = 0
else:
bad_epochs += 1
if bad_epochs >= patience:
print("Early stopping.")
break
model.load_state_dict(best_state)
print("Best val MAE:", best_val)
model.eval()
test_preds = []
with torch.no_grad():
for x, lengths in test_loader:
x = x.to(device)
lengths = lengths.to(device)
pred = model(x, lengths)
pred = denormalize_y(pred.cpu().numpy())
test_preds.append(pred)
test_preds = np.concatenate(test_preds)
test_preds = np.rint(test_preds).astype(np.int64)
submission = pd.DataFrame({
"ID": test_df["ID"],
"output": test_preds
})
submission.to_csv("submission.csv", index=False)
submission.head()
ID output
0 0 6
1 1 -9
2 2 7
3 3 3
4 4 33
import pandas as pd
from sklearn.metrics import mean_absolute_error
try:
submission_df = pd.read_csv('/content/submission.csv')
solution_df = pd.read_csv('/content/solution.csv')
comparison_df = pd.merge(submission_df, solution_df, on='ID', suffixes=('_predicted', '_actual'))
mae = mean_absolute_error(comparison_df['output_actual'], comparison_df['output_predicted'])
print(f"\nMean Absolute Error (MAE) on the test set: {mae:.4f}")
print("\nComparison of predicted vs. actual outputs (first 5 rows):")
display(comparison_df.head())
except FileNotFoundError:
print("\nError: 'solution.csv' not found. Please ensure the ground truth file is in the /content/ directory to evaluate the model.")
except Exception as e:
print(f"\nAn error occurred during MAE calculation: {e}")
Mean Absolute Error (MAE) on the test set: 21.9046
Comparison of predicted vs. actual outputs (first 5 rows):
ID output_predicted output_actual Usage
0 0 6 4 Public
1 1 -9 -6 Private
2 2 7 0 Public
3 3 3 0 Private
4 4 33 29 Private