Checklist IOAI TST 2025 IOAI Team Selection Test · Task 3
SOTA-AI Editorial
LeetCode Difficulty Prediction from Solution Descriptions
Predict whether a LeetCode problem is easy, medium or hard from a step-by-step natural-language description of its Python solution.
Hints
Hint 1
Start with a pretrained Transformer model and use it to extract embeddings from the code.
Hint 2
Normalize the embeddings before passing them to the classifier.
Hint 3
A lightweight ensemble model such as AdaBoost or Random Forest can work well on top of the Transformer embeddings.
The solution in one line
Use a pretrained RoBERTa model to turn each code sample into a meaningful embedding, normalize these embeddings, and classify them with AdaBoost.
The editorial
Looking at the data
This is a three-class classification task where the goal is to predict whether a code sample is easy, medium, or hard.
I started by loading the training and test data and splitting the training set into 80% for training and 20% for validation. I used a stratified split with random_state=42 so that the class distribution stays roughly the same in both sets.
For the code representation, I used the pretrained FacebookAI/roberta-base model. Missing code values were replaced with empty strings, and the code was tokenized with a maximum length of 256 tokens.
Rather than using the representation of a single token, I applied mean pooling to the RoBERTa hidden states. This gives one 768-dimensional vector for each code sample while ignoring padding tokens. I then standardized these embeddings with StandardScaler, fitting the scaler only on the training data.
Since there are three difficulty classes, I used macro F1 as the main validation metric. This gives the same importance to easy, medium, and hard, which is useful if the classes are not perfectly balanced.
The approach
I wanted to keep the Transformer part relatively simple and use it mainly as a feature extractor.
I first used RoBERTa to convert each code sample into a 768-dimensional embedding. The Transformer itself was kept frozen, so there was no fine-tuning. After standardizing the embeddings, I trained an AdaBoost classifier on top of them.
The final pipeline is:
- RoBERTa: Generate contextual embeddings from the code.
- Mean pooling: Turn the token-level representations into one vector per sample.
- Standardization: Normalize the embeddings using statistics from the training set.
- AdaBoost: Train 200 shallow decision trees with
max_depth=2,min_samples_leaf=5, and a learning rate of0.05. - Prediction: Convert the predicted class IDs back to
easy,medium, andhard.
I chose AdaBoost because the RoBERTa embeddings already contain a lot of information about the code, so a relatively lightweight classifier can be used on top of them. I also used balanced class weights in the decision trees to help with differences between the classes.
Results
| Method | Score |
|---|---|
| baseline | 0.26170 |
| RoBERTa + AdaBoost | 0.37509 |
| Author's solution | 0.42881 |
The RoBERTa + AdaBoost approach improved my original score from 0.26170 to 0.37509, which is a gain of 0.11339 F1. It still remains below the author's solution at 0.42881, so there is room for further improvement.
The reported 0.37509 is the score obtained on the released validation set, while 0.42881 is the author's competition result.
Running time
Most of the computation comes from extracting the RoBERTa embeddings. The code samples are processed in batches of 8 with a maximum sequence length of 256, so using a GPU makes this step considerably faster.
Once the embeddings have been extracted, training and prediction with AdaBoost are relatively lightweight compared with running the Transformer.
RoBERTa is used only in inference mode with torch.no_grad(), so no gradients are stored and the model is not fine-tuned. The embeddings are extracted once for the training, validation, and test sets and then reused by the classifier.
I did not record the exact execution time, so I am not giving a specific CPU or GPU runtime.
Code
Solution code
import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModel
from tqdm.auto import tqdm
from sklearn.ensemble import AdaBoostClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import f1_score, classification_report
from sklearn.model_selection import train_test_split
train = pd.read_csv("/kaggle/input/tst-day-3-upsolving/train.csv")
test = pd.read_csv("/kaggle/input/tst-day-3-upsolving/test.csv")
train_df, val_df = train_test_split(
train,
test_size=0.2,
random_state=42,
stratify=train["label"]
)
print("Train :", train_df.shape)
print("Validation :", val_df.shape)
MODEL_NAME = "FacebookAI/roberta-base"
MAX_LENGTH = 256
BATCH_SIZE = 8
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
encoder = AutoModel.from_pretrained(MODEL_NAME).to(device)
encoder.eval()
class CodeDataset(Dataset):
def __init__(self, texts, tokenizer, max_length=256):
self.texts = texts.fillna("").astype(str).tolist()
self.tokenizer = tokenizer
self.max_length = max_length
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
return self.tokenizer(
self.texts[idx],
truncation=True,
padding="max_length",
max_length=self.max_length,
return_tensors="pt"
)
def extract_embeddings(texts):
dataset = CodeDataset(texts, tokenizer, MAX_LENGTH)
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False)
all_embeddings = []
with torch.no_grad():
for batch in tqdm(loader):
input_ids = batch["input_ids"].squeeze(1).to(device)
attention_mask = batch["attention_mask"].squeeze(1).to(device)
outputs = encoder(
input_ids=input_ids,
attention_mask=attention_mask
)
hidden = outputs.last_hidden_state
mask = attention_mask.unsqueeze(-1).float()
summed = (hidden * mask).sum(dim=1)
counts = mask.sum(dim=1).clamp(min=1)
embeddings = summed / counts
all_embeddings.append(embeddings.cpu().numpy())
return np.concatenate(all_embeddings, axis=0)
sc = StandardScaler()
X_train_emb = extract_embeddings(train_df["code"])
X_train_emb = sc.fit_transform(X_train_emb)
X_val_emb = extract_embeddings(val_df["code"])
X_val_emb = sc.transform(X_val_emb)
X_test_emb = extract_embeddings(test["code"])
X_test_emb = sc.transform(X_test_emb)
y_train = train_df["label"].to_numpy()
y_val = val_df["label"].to_numpy()
print("Train embeddings:", X_train_emb.shape)
print("Validation embeddings:", X_val_emb.shape)
print("Test embeddings:", X_test_emb.shape)
ada = AdaBoostClassifier(
estimator=DecisionTreeClassifier(
max_depth=2,
min_samples_leaf=5,
class_weight="balanced",
random_state=42
),
n_estimators=200,
learning_rate=0.05,
random_state=42
)
ada.fit(X_train_emb, y_train)
val_pred = ada.predict(X_val_emb)
print("Macro F1:", f1_score(
y_val,
val_pred,
average="macro"
))
print(classification_report(
y_val,
val_pred,
labels=[0, 1, 2],
target_names=["easy", "medium", "hard"],
zero_division=0
))
test_pred = ada.predict(X_test_emb).astype(int)
id2label = {
0: "easy",
1: "medium",
2: "hard"
}
submission = pd.DataFrame({
"id": test["id"],
"difficulty": [id2label[int(p)] for p in test_pred]
})
submission.to_csv(
"submission_roberta_ada.csv",
index=False
)Editorials are written by members of the SOTA community, not by the olympiad. The official statement and scoring win wherever an editorial differs.