IOAI Training: Learn AI by solving problems

The problem in brief

You get 6 606 Peruvian court rulings in which every sensitive name was replaced by the token xxxxxxxxxx (23 131 replacements, 2 to 5 per document). You also get a "leak": for every document, the embedding of the original text (with the real names) and the embedding of the redacted text. A document embedding is the average of the unit-normalised Spanish fastText vectors of its words. Finally you get 800 candidate vectors: 300 are noisy copies of the true name vectors, 500 are decoys, all shuffled. For every xxxxxxxxxx you must output the hidden name as a word.

There are no labels at all, neither on Kaggle's side for us (test answers are private) nor for training. We never submitted anything. All scores below are self-consistency measurements, explained in the Results section.

How you are scored

\text{Score} = 0.7\,\text{bag} + 0.3\,\text{positional},

where bag is the fraction of names recovered per document regardless of order (multiset intersection), and positional is the fraction of slots with exactly the right name.

What to aim for. The only public-leaderboard entry is the organisers' baseline, 0.211.

First look at the data

Two facts from the statement are the whole key. The embedding of a document with $n$ words is

\text{full} = \frac{1}{n}\Big(\sum_{\text{other words}} u_w + \sum_{j=1}^{k} u_{\text{name}_j}\Big), \qquad \text{redacted} = \frac{1}{n}\Big(\sum_{\text{other words}} u_w + k\, e_{\text{bot}}\Big),

with $u$ the unit word vectors and $e_{\text{bot}}$ the vector of xxxxxxxxxx. Everything that is not a name cancels in the difference, so

s_d = n\,(\text{full} - \text{redacted}) + k\, e_{\text{bot}} = \sum_{j=1}^{k} u_{\text{name}_j}.

Every document gives us the exact sum of its $k$ hidden name vectors. The task becomes: write 6 606 known vectors as sums of $k$ unknown "atoms" from a common dictionary of 300 names.

The candidates make this harder than it looks. Name vectors in fastText are strongly correlated (surnames live in one region of the space: the median cosine between two candidates is 0.30, the top 1 % above 0.73), and the candidates are noisy: once we recovered the true vectors, their cosine with the given candidates was only about 0.8.

A first baseline

Decode each $s_d$ with the 800 given candidates: greedily add the candidate that most reduces $\|s_d - \text{sum}\|^2$ , $k$ times, then repeatedly try swapping each chosen atom for a better one. With the raw candidates the median residual is 1.09: no document is explained exactly, because the noise on each candidate adds up.

The key idea

Use the 6 606 documents to clean the dictionary. This is dictionary learning by alternating minimisation:

Decode every document with the current atoms (greedy + swaps).
Re-estimate every atom by least squares. Let $M_{dj}$ be how many times atom $j$ was chosen in document $d$ ; then $S \approx M V$ , and $V = \arg\min \|S - MV\|^2$ pools the evidence of every document that uses a name. Since each $s_d$ is an exact sum, a name used in a hundred documents is pinned down very precisely.
Renormalise each atom to unit length (true word vectors have norm 1) and drop atoms that no document uses. Decoys are rarely the best choice once the true names sharpen, so they die out.

Alternating decoding and refitting, final residuals, and the match with fastText words

After four rounds, 86 % of the documents are explained exactly (residual below 0.05, against a typical norm of 2 to 4 for $s_d$ ); the process stabilises at 90 % of the documents, 88.5 % of the names. The right panel shows the strongest confirmation that the atoms are right: we streamed the 1.5 million alphabetic words of the fastText Spanish vocabulary and computed each vector's closest word. The given candidates match no word well (median cosine 0.77), while the re-estimated atoms match a word almost perfectly (median cosine 0.9998): we have recovered the exact vectors of names such as castro, chavez, trujillo, vargas.

Some atoms stay far from every vocabulary word (cosine 0.5 to 0.75). One of them is used 429 times. Its closest words are sequeira and jsequeiros: it is very likely the surname Sequeiros, whose lowercase form is not in the .vec vocabulary. fastText builds such vectors from character n-grams, which needs the full .bin model (7 GB) that we did not download. These names get the nearest vocabulary word, which is probably wrong.

The order is a second, separate problem. An average is blind to word order, so the embeddings say nothing about which name sits in which slot. Choosing a random order would give about 36 % positional agreement on our decoded bags. The text helps: legal documents repeat formulas ("el señor juez xxxxxxxxxx Saldarriaga", "Barrios xxxxxxxxxx") and the same people appear in many rulings. We fit a naive-Bayes context model $P(\text{neighbour word} \mid \text{name})$ on the previous and next one and two words, by hard EM: start with every bag name at every slot, fit the counts, give each document the permutation of its bag with the highest likelihood, refit, repeat. It converges in three rounds.

The full solution

Compute $s_d$ for every document (float64).
Start from all 800 candidates; 25 rounds of decode, least-squares refit, unit renormalisation and pruning; a final decoding with 10 swap sweeps. 435 atoms survive, some of them near-duplicates of the same word (four atoms map to castro), which is harmless because we merge them by name.
Name each atom by its closest fastText word (lowercased, accents removed, as the scorer does).
Order the names inside each document with the EM context model.

A, res = decode(S, k, V)                      # which atoms, for every document
V = refit(S, A, len(V))[use >= 2]             # least squares + unit norm + drop unused atoms

Results

No ground truth is available, so we measure what we can.

quantity	value
organisers' baseline (public leaderboard)	0.211
documents explained exactly with the raw candidates	0 %
documents explained exactly after dictionary learning	90.0 %
names in exactly explained documents	88.5 %
of which the atom matches a vocabulary word (cosine > 0.99)	79.8 % of all names
order: agreement with unredacted text evidence, EM context model	79.9 % (11 480 slots)
order: same check for a random order	33.7 %

The bag score should lie between 0.80 (names that are both exactly reconstructed and exactly named) and about 0.88 (all names in exactly explained documents, if most out-of-vocabulary names still happen to be right). For the order, we used an independent check: in about half of the slots, the document itself contains the same neighbour word next to one of the bag names in clear text (for example "Prado Saldarriaga" written out elsewhere in the ruling). On those slots our EM order agrees with that evidence 80 % of the time, against 34 % for a random order. These slots are easier than average, so the true positional accuracy is probably lower. Our rough estimate of the final score is therefore around $0.7 \times 0.84 + 0.3 \times 0.65 \approx 0.78$ , nearly four times the baseline; this is an estimate, not a measured score.

What to remember

Write down the formula of the embedding. The averaging made everything except the names cancel exactly. Algebra on the representation beat any learned model.
An exact signal plus many observations means you can learn the dictionary. Least squares over thousands of documents turns noisy candidates into exact vectors; this is the alternating idea behind sparse coding, k-means and matrix factorisation.
Use physical constraints. Unit norm for word vectors and integer counts for names kept the alternation stable.
Check your answers against an independent source. A near-perfect match with fastText vocabulary words, and agreement with unredacted text, are strong evidence even without labels.
Averaged embeddings leak. Removing a name from a document does not remove it from an embedding computed before redaction: a privacy lesson as much as an ML one.

Note on the code. The full solution.py below reproduces the setup that produced the reported score, but it was not re-run end to end after it was assembled.

"""Embedding Secrets (Peru OPEIA 2026, NLP): recover redacted names from leaked document embeddings.

1. Algebra: a document embedding is the MEAN of unit word vectors, so
       s_d = n_words * (full_emb - redacted_emb) + n_redacted * e_bot
   is exactly the SUM of the unit vectors of the hidden names of document d.
2. Sparse decoding: write each s_d as a sum of n_redacted atoms (greedy + swap search).
3. Dictionary learning: re-estimate every atom by least squares from all documents that use
   it, renormalise to unit norm, drop unused atoms, repeat. The noisy candidates converge to
   the exact word vectors, and the decoys disappear.
4. Naming: stream the fastText cc.es.300.vec vocabulary and give each atom its closest word.
5. Order: the embedding says nothing about order, so a context model (previous/next words,
   hard EM over the permutations of each bag) puts each name at its most likely slot.

Data: kaggle competition "ioai-nlp-challenge-2026"; fastText vectors from
      https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.es.300.vec.gz (streamed, not stored).
Usage: curl -s <vec.gz url> | python solution.py     -> submission.csv (not submitted)
Runtime: ~2 min of decoding + ~15 min to stream the vocabulary; < 2 GB RAM.
"""
import collections, gzip, itertools, re, sys, unicodedata
import numpy as np, pandas as pd


def load():
    d = pd.read_parquet('challenge_dataset.parquet')
    C = np.load('candidates.npy').astype(np.float64); e = np.load('e_bot.npy').astype(np.float64)
    F = np.stack(d.full_emb.values).astype(np.float64); R = np.stack(d.redacted_emb.values).astype(np.float64)
    k = d.n_redacted.values
    return d, C, d.n_words.values[:, None] * (F - R) + k[:, None] * e, k


def decode(S, k, V, sweeps=5):
    """Multiset of k atoms per document whose sum is closest to s_d: greedy, then swap each slot."""
    A = np.full((len(S), k.max()), -1); nrm = (V ** 2).sum(1); R = S.copy()
    for t in range(k.max()):
        act = k > t
        j = (R[act] @ V.T - 0.5 * nrm).argmax(1); A[act, t] = j; R[act] -= V[j]
    for _ in range(sweeps):
        for t in range(k.max()):
            act = np.where(k > t)[0]
            Rp = R[act] + V[A[act, t]]
            j = (Rp @ V.T - 0.5 * nrm).argmax(1); A[act, t] = j; R[act] = Rp - V[j]
    return A, np.linalg.norm(R, axis=1)


def refit(S, A, m):
    M = np.zeros((len(S), m))                              # M[d, j] = how many times atom j is in doc d
    for t in range(A.shape[1]):
        ok = A[:, t] >= 0; np.add.at(M, (np.where(ok)[0], A[ok, t]), 1)
    V = np.linalg.lstsq(M, S, rcond=None)[0]
    return V / np.maximum(np.linalg.norm(V, axis=1, keepdims=True), 1e-12)   # word vectors have unit norm


def nearest_words(V, stream, chunk=20000):
    best, bw = np.full(len(V), -2.0), np.empty(len(V), object)
    f = gzip.open(stream, 'rt', encoding='utf-8', errors='ignore'); next(f)
    words, vecs = [], []
    def flush():
        W = np.array(vecs, np.float32); W /= np.linalg.norm(W, axis=1, keepdims=True) + 1e-9
        sim = V @ W.T.astype(np.float64); j = sim.argmax(1); upd = sim[np.arange(len(V)), j] > best
        best[upd] = sim[np.arange(len(V)), j][upd]; bw[upd] = np.array(words, object)[j[upd]]
        words.clear(); vecs.clear()
    for line in f:
        p = line.rstrip().split(' ')
        if len(p) == 301 and p[0].isalpha():
            words.append(p[0]); vecs.append(np.asarray(p[1:], np.float32))
            if len(words) == chunk:
                flush()
    if words:
        flush()
    return bw, best


norm = lambda t: unicodedata.normalize('NFKD', t).encode('ascii', 'ignore').decode().lower()


def contexts(text):
    t = [norm(x) for x in re.findall(r'\w+|[^\w\s]', text)]
    g = lambda i: t[i] if 0 <= i < len(t) else '<s>'
    return [[('p1', g(i - 1)), ('n1', g(i + 1)), ('p2', g(i - 2) + '_' + g(i - 1)), ('n2', g(i + 1) + '_' + g(i + 2))]
            for i, x in enumerate(t) if x == 'xxxxxxxxxx']


def order(ctx, bags, iters=4):
    """Hard EM: fit P(context word | name) from the current assignment, re-assign each bag optimally."""
    assign = None
    for it in range(iters):
        cnt, tot = collections.Counter(), collections.Counter()
        for c, b, P in zip(ctx, bags, assign or [None] * len(bags)):
            for slot, feats in enumerate(c):
                names = b if P is None else [P[slot]]          # first round: every bag name at every slot
                for n in names:
                    for f in feats:
                        cnt[(f, n)] += 1 / len(names)
                    tot[n] += 1 / len(names)
        lp = lambda f, n: np.log((cnt[(f, n)] + 0.1) / (tot[n] + 500.0))
        assign = [max(set(itertools.permutations(b)), key=lambda p: sum(lp(f, n) for fs, n in zip(c, p) for f in fs))
                  for c, b in zip(ctx, bags)]
    return assign


if __name__ == '__main__':
    d, C, S, k = load()
    V = C.copy()
    for it in range(25):
        A, res = decode(S, k, V)
        use = np.bincount(A[A >= 0], minlength=len(V))
        V = refit(S, A, len(V))[use >= 2]
        print(f'iteration {it}: {len(V)} atoms, documents explained exactly: {np.mean(res < 0.05):.3f}', flush=True)
    A, res = decode(S, k, V, sweeps=10)
    words, sim = nearest_words(V, sys.stdin.buffer)
    name = np.array([norm(w) for w in words])
    print(f'atoms matching a vocabulary word (cos > 0.99): {np.mean(sim > 0.99):.3f}')
    bags = [[name[a] for a in row if a >= 0] for row in A]
    assign = order([contexts(t) for t in d.redacted_text], bags)
    rows = [(f'{i}_{j}', n) for i, P in enumerate(assign) for j, n in enumerate(P)]
    pd.DataFrame(rows, columns=['Id', 'predicted_name']).to_csv('submission.csv', index=False)
    print('wrote submission.csv (not submitted)')