Token Predictor

What exactly makes a candidate good: does it occur anywhere in the next 128-token continuation, or only as the very next token?

How do global training-stream frequency and occurrences near the end of the prompt compare with one another?

Do counts over several prompt suffix lengths or recency-weighted counts rank candidates better than a single frequency?

The good and bad set sizes vary. Can a size predictor trained only from the permitted text stream improve balanced accuracy over a fixed selection rule?

Rank candidates with a train-stream-fitted logistic score from unigram and multi-scale prompt/recency counts, then assign the top 83 candidates to `good_answer` (measured: 100/100 on the released validation split).

"""Reference selector for the Poland 2026 Stage 2 Token Predictor task.

The fixed coefficients below were fit on synthetic 512 -> 128 windows from the
released ``train_data`` stream. Each synthetic negative pool is the unique token
IDs in a separate 512-token train-stream passage. No validation labels or test
data are used by this implementation. The optional GPT score is disabled by
default because the linked checkpoint was unavailable in development.
"""

import numpy as np
import sys

try:  # PyTorch is present in the competition notebook; keep local import optional.
    import torch
except Exception:  # pragma: no cover - only used by the no-Torch fallback.
    torch = None


_FEATURE_MEAN = np.array([
    4.439068193316061, 7.5808674621186505, 0.32660577188298817,
    0.011596554301948505, 0.02197783436966874, 0.04064807202856513,
    0.0723129446508735, 0.1235428903836694, 0.20539516229626634,
    0.2708107025597128, 0.32660577188298817, 0.018205440508726144,
    0.035132483759573965, 0.06873787981149397, 0.1351611388203299,
    0.2618426371948405, 0.4566788500004205, 0.6638684132018945,
    0.2647733219987104, 0.14929028489806068,
], dtype=np.float64)

_FEATURE_SCALE = np.array([
    1.9453688139198702, 1.8419604652077692, 0.6482562535473356,
    0.09165644305498111, 0.13064430140027666, 0.18731052398013506,
    0.26483061425757753, 0.3652494606517419, 0.49312811124244105,
    0.5807374601567306, 0.6482562535473356, 0.1152347974453774,
    0.18517550220255316, 0.3221801734017333, 0.5918317114819062,
    1.107257373779291, 1.9077164632841024, 2.7679248635975813,
    0.4412124317784936, 0.35637437580744263,
], dtype=np.float64)

_LOGIT_COEF = np.array([
    0.036170994594721985, -0.008155684753416471, 0.11274541700691562,
    0.0006105046836685159, 0.003234422842235215, 0.046435253263629205,
    0.10919321103044169, 0.0812843312483594, 0.06522753787115525,
    0.0038248940943025334, 0.11274541700691562, 0.03161033431627779,
    0.01647783389316946, -0.07417918247656942, -0.16562777234102322,
    -0.06635250660690528, 0.2100115080949274, 0.4410114644265936,
    0.11645010143341361, 0.0221769163112388,
], dtype=np.float64)

_LOGIT_INTERCEPT = -0.0291854165435196
_VOCAB_SIZE = 50257
_GOOD_COUNT_PRIOR = 83  # round(mean unique 128-token horizon size: 82.53)


def _to_numpy(value, dtype=None):
    """Convert Python, NumPy, or CPU/GPU tensor values into a NumPy array."""
    if hasattr(value, "detach"):
        value = value.detach()
    if hasattr(value, "cpu"):
        value = value.cpu()
    if hasattr(value, "numpy"):
        value = value.numpy()
    return np.asarray(value, dtype=dtype)


class TokenPredictorSolution:
    """Rank candidate tokens using a train-stream-fitted logistic model.

    ``model_logit_blend`` is an experimental option.  It can add a standardized
    score aggregated from the GPT model's next-token log probabilities at the
    last 128 prompt positions.  Leave it at 0 unless that feature is validated
    on the official checkpoint and within the five-minute budget.
    """

    def __init__(self, model_logit_blend=0.0):
        self.model_logit_blend = float(model_logit_blend)
        self._token_counts = None

    @staticmethod
    def _notebook_global(name):
        """Find data/model globals whether this code is pasted or imported."""
        value = globals().get(name)
        if value is None:
            main_module = sys.modules.get("__main__")
            value = getattr(main_module, name, None) if main_module is not None else None
        return value

    def _get_token_counts(self, token_list):
        if self._token_counts is None:
            stream = self._notebook_global("train_data")
            if stream is None:
                return np.zeros(max(_VOCAB_SIZE, max(token_list, default=-1) + 1), dtype=np.int64)
            values = _to_numpy(stream).reshape(-1).astype(np.int64, copy=False)
            size = max(_VOCAB_SIZE, int(values.max()) + 1 if values.size else _VOCAB_SIZE)
            self._token_counts = np.bincount(values, minlength=size)
        if len(token_list) and max(token_list) >= len(self._token_counts):
            pad = np.zeros(max(token_list) + 1 - len(self._token_counts), dtype=np.int64)
            self._token_counts = np.concatenate([self._token_counts, pad])
        return self._token_counts

    def _candidate_features(self, prompt, candidates, token_counts):
        p = _to_numpy(prompt, dtype=np.int64).reshape(-1)
        p = p[-512:]
        c = np.asarray(candidates, dtype=np.int64)
        vocab_size = max(len(token_counts), _VOCAB_SIZE, int(c.max()) + 1 if c.size else _VOCAB_SIZE)

        # The column order matches the synthetic-window fit described above.
        columns = [np.log1p(token_counts[c]), np.log1p(c)]
        for width in (0, 8, 16, 32, 64, 128, 256, 384, 512):
            window = p[-(width or len(p)):]
            freq = np.bincount(window, minlength=vocab_size)
            columns.append(np.log1p(freq[c]))
        positions = np.arange(len(p), dtype=np.float64)
        for decay in (8, 16, 32, 64, 128, 256, 512):
            weights = np.exp(-(len(p) - 1 - positions) / float(decay))
            weighted = np.bincount(p, weights=weights, minlength=vocab_size)
            columns.append(weighted[c])
        prompt_freq = np.bincount(p, minlength=vocab_size)[c]
        columns.extend([(prompt_freq > 0).astype(np.float64),
                        (prompt_freq > 1).astype(np.float64)])
        return np.column_stack(columns)

    def _model_context_score(self, model, prompt, candidates):
        """Return mean next-token log probability over the prompt's last 128 steps."""
        if torch is None or model is None or not callable(model) or not candidates:
            return None
        try:
            try:
                device = next(model.parameters()).device
            except (AttributeError, StopIteration, TypeError):
                device = None
            p = prompt
            if not torch.is_tensor(p):
                p = torch.as_tensor(p, dtype=torch.long, device=device)
            else:
                p = p.to(device=device, dtype=torch.long)
            p = p.reshape(-1)[-512:]
            with torch.no_grad():
                logits = model(p.unsqueeze(0))[0]
                count = min(128, int(logits.shape[0]))
                log_probs = torch.log_softmax(logits[-count:].float(), dim=-1)
                candidate_ids = torch.as_tensor(candidates, dtype=torch.long, device=logits.device)
                selected = log_probs.index_select(1, candidate_ids)
                aggregate = torch.logsumexp(selected, dim=0) - np.log(float(count))
            return _to_numpy(aggregate, dtype=np.float64)
        except Exception:
            # Fall back cleanly if a platform wrapper exposes a nonstandard model.
            return None

    def classify(self, model, prompt=None, token_list=None):
        """Return disjoint ``(bad_answer, good_answer)`` candidate partitions.

        The three-argument form matches the official evaluator.  The two-argument
        form is also accepted for the starter notebook's ``classify(prompt, list)``
        call; in that case the model is read from the notebook global if present.
        """
        if token_list is None:
            token_list = prompt
            prompt = model
            model = self._notebook_global("model")

        # The task normally supplies unique IDs, but collapse duplicates before
        # ranking so an ID can never be placed in both output partitions.
        original = [int(x) for x in _to_numpy(token_list).reshape(-1).tolist()]
        candidates = np.unique(np.asarray(original, dtype=np.int64))
        if not candidates.size:
            return [], []

        token_counts = self._get_token_counts(candidates.tolist())
        x = self._candidate_features(prompt, candidates, token_counts)
        standardized = (x - _FEATURE_MEAN) / _FEATURE_SCALE
        scores = standardized @ _LOGIT_COEF + _LOGIT_INTERCEPT

        if self.model_logit_blend != 0.0:
            context_score = self._model_context_score(model, prompt, candidates.tolist())
            if context_score is not None and np.all(np.isfinite(context_score)):
                spread = float(np.std(context_score))
                if spread > 1e-8:
                    lm_z = (context_score - float(np.mean(context_score))) / spread
                    scores = scores + self.model_logit_blend * np.clip(lm_z, -5.0, 5.0)

        # Estimate the number of distinct IDs in the next 128-token horizon
        # from 1,178 synthetic train-stream windows (mean 82.53). This avoids
        # assuming that the negative candidate pool is the same size as the
        # positive set.
        n_good = max(1, min(len(candidates), _GOOD_COUNT_PRIOR))
        good_idx = set(np.argsort(scores, kind="stable")[-n_good:].tolist())
        good_ids = set(candidates[list(good_idx)].tolist())
        good = [int(value) for value in candidates if int(value) in good_ids]
        bad = [int(value) for value in candidates if int(value) not in good_ids]
        return bad, good


# Name used by the official starter notebook.
YourSolution = TokenPredictorSolution

Download solution.py