Discord

Checklist OAI 2026 Stage II · Task 2

Token Predictor

Polish title: Predyktor Tokenów

Using GPT-2 XL and a 512-token Shakespeare prompt, split 256 candidate tokens into those that do and do not occur among the next 128 tokens.

  • NLP
  • Token set classification with a pre-trained language model
  • Polish original · English translation
Download the English notebook Open the task notebook (Polish)

The task

Large language models such as GPT-2 assign probabilities to possible next tokens, reflecting statistical knowledge of text. The contestant has a pre-trained GPT-2 XL model (standard OpenAI weights) and text made of verses from various plays by William Shakespeare, already tokenised with the GPT-2 BPE encoding.

Given a 512-token prompt and a token_list of exactly 256 tokens, the solution must decide which tokens appear in the next q_size = 128 tokens of the original text (good_tokens) and which do not (bad_tokens); the list contains 128 of each in random order. YourSolution.classify(model, prompt, token_list) returns bad_answer and good_answer, which must form a disjoint partition of token_list. The model definition may be modified (for example, a model.query(prompt, token_list) method returning logits is provided).

The training data are one contiguous text of 301,966 tokens; the validation set has 99 examples built from another fragment of the same text, and the hidden test set has 100 examples.

Abridged and translated by SOTA from the official Polish materials. The official statement has the exact rules, and it wins wherever this summary differs.

In English

This task was published in Polish. SOTA translated it into English on 16 September 2026. Only the words changed in the notebooks: markdown, code comments, messages and printed output. The code, file names and paths are the original's, so a translated notebook runs with the original data.

Read the task notebook in English 24506 words and 8 code cells

Token Predictor

image.png

Image generated with DeepAI.

Introduction

Large language models (LLMs), such as GPT-2, learn statistical relationships between tokens in text. At each step, the model processes a sequence of tokens and assigns a certain probability to every possible next token. This probability distribution reflects the model's knowledge of which words and phrases naturally follow one another in a given context.

Task

You have at your disposal a pre-trained GPT-2 XL model and a text dataset containing lines taken from various plays by William Shakespeare. Your task is to build the best possible model which, on the basis of a given text fragment (of size 512 tokens), is able to predict which tokens from a list will appear later in the text (within the next q_size = 128 tokens).

Data

The dataset consists of three parts:

  1. The training dataset, which is a continuous fragment of text consisting of 301966 tokens. The data come from lines of various plays by William Shakespeare.

  2. The validation dataset - 99 examples (built from a different fragment of the same text as the training set), each consisting of:

  • prompt — a sequence of 512 context tokens.
  • good_tokens - q_size = 128 tokens that actually appear after the prompt in the original text.
  • bad_tokens — q_size = 128 tokens that do not appear in the next q_size tokens after the prompt (negatives).
  • token_list - contains exactly 256 tokens, of which 128 are good_tokens and 128 are bad_tokens, in random order.
  1. The test dataset - 100 examples that will be used to score your solution. You do not have access to the test data while solving the task.

Scoring Criterion

Given the prompt and the token_list, your solution must distinguish the good_tokens from the bad_tokens and return good_answer and bad_answer, which are, respectively, the tokens predicted by the model as belonging to each of these groups. The two lists should form a disjoint partition of token_list, i.e. each token from token_list must be assigned to exactly one of them.

The quality of your solution will be measured by classification accuracy:

score=12∗(∣set(good_tokens)∩set(good_answer)∣∣set(good_tokens)∣+∣set(bad_tokens)∩set(bad_answer)∣∣set(bad_tokens)∣)\text{score} = \frac{1}{2}*(\frac{|\text{set}(good\_tokens) \cap \text{set}(good\_answer)|}{|\text{set}(good\_tokens)|} + \frac{|\text{set}(bad\_tokens) \cap \text{set}(bad\_answer)|}{|\text{set}(bad\_tokens)|})

You can score from 0 to 100 points for this task. The score will be scaled linearly depending on the mean accuracy achieved on the whole test set:

  • Accuracy ≤ 0.62: 0 points
  • Accuracy ≥ 0.68: 100 points
  • Values in between: scaled linearly.

The final score is calculated according to the formula:

Points={0for Accuracy≤0.62100×Accuracy−0.620.68−0.62for 0.62<Accuracy<0.68100for Accuracy≥0.68\text{Points} = \begin{cases} 0 & \text{for } \text{Accuracy} \leq 0.62 \\ 100 \times \frac{\text{Accuracy} - 0.62}{0.68 - 0.62} & \text{for } 0.62 < \text{Accuracy} < 0.68 \\ 100 & \text{for } \text{Accuracy} \geq 0.68 \end{cases}

Constraints

  • Your solution will be tested on the Contest Platform without internet access.
  • The evaluation of your final solution must not take longer than 5 minutes (with a GPU).
  • List of permitted libraries: numpy, pytorch.

Submission Files

This notebook, completed with your solution (see the class YourSolution).

Evaluation

Remember that during grading the FINAL_EVALUATION_MODE flag will be set to True.

You can score between 0 and 100 points for this task. The number of points you receive will be calculated on the (secret) test set on the Contest Platform according to the formula given above, rounded to an integer. If your solution does not meet the criteria above or does not run correctly, you will receive 0 points for the task.

Starter Code

In this section we initialise the environment by importing the required libraries and functions, and we load the GPT-2 model. The prepared code will make it easier for you to work with the data efficiently and to build a proper solution.

######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################
FINAL_EVALUATION_MODE = False # During grading on the grading system, the flag is automatically changed to True.
# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import numpy as np
import os
import random

# Setting the seed of the pseudo-random number generator to ensure deterministic results
seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)

device      = 'cuda' if torch.cuda.is_available() else 'cpu'
block_size  = 1024
prompt_size = 512
q_size      = 128   # must be equal to WINDOW_LEN in AuthorsSolution
assert prompt_size + 3 * q_size <= block_size

print(f"Device in use: {device}")

Model Definition

The GPT-2 model used in the task is defined below. You may change this code freely - e.g. add new methods. For example, we have created a method model.query(prompt, token_list) that returns logits on the basis of the prompt - described below.

class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=True)
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=True)
        self.attn_dropout  = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.n_head  = config.n_head
        self.n_embd  = config.n_embd
        self.dropout = config.dropout

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=None,
                dropout_p=self.dropout if self.training else 0, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_dropout(self.c_proj(y))


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.gelu    = nn.GELU()
        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp  = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class GPT(nn.Module):
    def __init__(self, config):
        super().__init__()
        from types import SimpleNamespace
        config = SimpleNamespace(**config)
        config.dropout = 0.0
        self.config = config
        self.transformer = nn.ModuleDict(dict(
            wte  = nn.Embedding(config.vocab_size, config.n_embd),
            wpe  = nn.Embedding(config.block_size, config.n_embd),
            drop = nn.Dropout(config.dropout),
            h    = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f = nn.LayerNorm(config.n_embd),
        ))
        print("Number of model parameters: %.2fM" % (sum(p.numel() for p in self.parameters()) / 1e6,))

    def forward(self, idx):
        b, t = idx.size()
        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)
        x = self.transformer.drop(self.transformer.wte(idx) + self.transformer.wpe(pos))
        for block in self.transformer.h:
            x = block(x)
        x = self.transformer.ln_f(x)
        return x @ self.transformer.wte.weight.mT

    def query(self, prompt, token_list):
        """
        For a given prompt and list of tokens, returns a list of logits:

        Args:
            prompt     (torch.Tensor): 1D tensor of context tokens, shape [T].
            token_list (list[int]):    list of tokens to classify.

        Returns:
            (list[int]): list of logits for the given tokens.
        """
        ctx = torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16) if device == 'cuda' else torch.amp.autocast(device_type='cpu', enabled=False)
        with torch.no_grad():
            with ctx:
                L = len(token_list)
                logits = self(prompt[None, ...])[0, -1, :]
                logits = logits.to(torch.float32).cpu().detach().numpy()[token_list].tolist()
                return logits

Loading the Model and the Data

The cell below loads the pre-trained GPT-2 XL model with standard weights, as well as the text data. The data are already tokenised — you do not need to install any tokeniser.

Available files:

  • ckpt.pt — GPT-2 XL weights (standard OpenAI weights, without any additional training),
  • train_data — the tokens of the training set (uint16, GPT-2 BPE encoding),
  • val_data — tuples (prompt, good_tokens, bad_tokens, token_list) used to evaluate the solution. You may assume that the test data will be in the same format as the validation data.
# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

weights_file = "ckpt.pt"
if not os.path.exists(weights_file):
    import gdown

    print("Downloading the model...")
    url = "https://drive.google.com/file/d/1aGenswqFszSBWYlXN26fkwpdyluTxfUI/view?usp=drive_link"
    gdown.download(url, weights_file, fuzzy=True)

print("Loading the pre-trained GPT-2 XL from ckpt.pt ...")
checkpoint = torch.load(weights_file, map_location=device, weights_only=False)
gptconf    = checkpoint['model_args']
model      = GPT(gptconf)
model.load_state_dict(checkpoint['model'])
model.to(device)
model.eval()
vocab_size = gptconf['vocab_size']
checkpoint = None
print("Model loaded successfully.")

train_data = torch.load("./data/train_data.pt", weights_only=False)
val_data   = torch.load("./data/val_data.pt", weights_only=False)

Data inspection

if not FINAL_EVALUATION_MODE:
    from tokenizers import Tokenizer

    enc = Tokenizer.from_file("./tokenizer.json") # we use the tokeniser to inspect the data - the data are already tokenised - you do not need to modify or use the tokeniser
    print("First 10 training tokens:", train_data[:10])
    print("First 10 training tokens as text:", enc.decode(train_data[:10]))

    for prompt, good_tokens, bad_tokens, token_list in val_data:
        print("First example from the validation set")
        print("Prompt (first 10 tokens):", prompt.tolist()[:10])
        print("Prompt (first 10 tokens) as text:", enc.decode(prompt.tolist()))
        print("Good tokens (first 10):", good_tokens[:10])
        print("Good tokens (first 10) as text:", [enc.decode([x]) for x in good_tokens])
        print("Bad tokens (first 10):", bad_tokens[:10])
        print("Bad tokens (first 10) as text:", [enc.decode([x]) for x in bad_tokens])
        print("All tokens (bad and good) (first 10):", token_list[:10])
        print("All tokens (bad and good) (first 10) as text:", [enc.decode([x]) for x in token_list])
        break

Code with the Scoring Criterion

# ######################### DO NOT CHANGE THIS CELL WHEN SUBMITTING ##########################

def accuracy(predicted, true_tokens):
    return len(set(predicted) & set(true_tokens)) / len(set(true_tokens))


def accuracy_to_score(acc: float) -> float:
    lo, hi = 0.62, 0.68
    if acc <= lo: return 0.0
    if acc >= hi: return 100.0
    return 100.0 * (acc - lo) / (hi - lo)


def evaluate_solution(solution, val_data) -> int:
    model.eval()
    accuracies = []
    for prompt, good_tokens, bad_tokens, token_list in val_data:
        # Classification
        token_list = token_list.copy()
        random.shuffle(token_list)
        bad_answer, good_answer = solution.classify(prompt, token_list)
        if bool(set(good_answer) & set(bad_answer)):
            accuracies.append(0)
            print("good_answer and bad_answer are not disjoint!")
            continue

        # We compute the accuracy for this example
        acc_bad = accuracy(bad_answer, bad_tokens)
        acc_good = accuracy(good_answer, good_tokens)

        # Mean of both groups for this example
        accuracies.append((acc_bad + acc_good) / 2)

    # Mean accuracy over all examples
    mean_acc = sum(accuracies) / len(accuracies)
    pts = int(round(accuracy_to_score(mean_acc)))

    print(f"\nFinal result:")
    print(f"  Analysed:           {len(val_data)} examples")
    print(f"  Mean accuracy:      {mean_acc:.4f}")
    print(f"  Points:             {pts} / 100")

    return pts

Your Solution

Place your solution in this section. You may also make changes to the model definition.

Your class YourSolution should implement the method classify(model, prompt, token_list), which returns bad_answer, good_answer:

  • bad_answer — the tokens classified as bad,
  • good_answer — the tokens classified as good.

The lists must be disjoint.

class YourSolution():
    def classify(self, prompt, token_list):
        L = len(token_list)
        return token_list[:L // 2], token_list[L // 2:]

Evaluation

Running the cell below lets you check how many points your solution would score on the validation data. Before submitting, make sure that the whole notebook runs from start to finish without errors and without any user intervention after selecting the "Run All" option.


if not FINAL_EVALUATION_MODE:
    your_solution = YourSolution()
    print(f"Score for YourSolution: {evaluate_solution(your_solution, val_data)} pts")

Translated by SOTA. The Polish original is the official version and wins wherever the two differ. The data (data/) and tokenizer.json sit next to the original notebook, and the notebook downloads the GPT-2 XL weights itself. As in the original, the text calls the method classify(model, prompt, token_list), but the code calls classify(prompt, token_list) and expects (bad_answer, good_answer) in that order. If you organise this olympiad and would like the translation removed, email [email protected] and we will take it down.

At a glance

You get
ckpt.pt (GPT-2 XL weights, Google Drive), data/train_data.pt and data/val_data.pt (tuples of prompt, good_tokens, bad_tokens, token_list) and tokenizer.json in the task folder.
You submit
This notebook with YourSolution.classify.
Scoring
Per example, score = ½·(|good ∩ good_answer|/|good| + |bad ∩ bad_answer|/|bad|), averaged over the test set. Points = 0 at ≤ 0.62, 100 at ≥ 0.68, 100 × (acc − 0.62)/(0.68 − 0.62) in between; rounded.
Rules
  • Tested without Internet access; evaluation must take at most 5 minutes (with a GPU).
  • Allowed libraries: numpy, pytorch.
Format
Stage II (regional, on site in Kraków, Poznań, Warsaw and Wrocław), 13–15 March 2026; two tasks per day in 5-hour sessions. Evaluated automatically on the Competition Platform (Platforma Konkursowa) on a hidden test set; points are rounded to an integer, and a notebook that fails the requirements or does not run scores 0.

In English (SOTA)

Notebook

Data

Models

Results

Source files

Details

Year
2026, Kraków, Poznań, Warsaw and Wrocław, Poland
Round
Stage II · Task 2
Language
Polish; English translation by SOTA
License
Not stated by the source