AICC (IOAI Community Contest) 2026 · Correction by IOAI Training

Tally

This page gives the full approach. Read it once you have searched.

Metric: micro F1 over complete names
Evaluated on: held-out validation (official test labels are private)
Our score: 0.699 entity F1 (2-fold on 50 gold)
Expected level: baseline 0.66, reference 0.83; LB best contestant 0.812

Our solution does not reach the expected level: there is still room to beat it

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

This is named-entity recognition (NER) with almost no labels. You must tag every word of 4 000 news sentences as part of a person, organisation, location, creative work, event, or nothing. Only 50 sentences are fully tagged. For 8 000 more you only know how many names of each type they contain, not where they are. The real names were replaced by invented words ("zhiggerry", "sovietnama") that never appear in two sentences, so memorising names is useless. The only pretrained model allowed is google/electra-base-discriminator.

The data are the official Kaggle files. The test labels are private and we did not submit anything: our scores come from our own validation on the 50 tagged sentences (2-fold) and on 500 held-out count sentences.

How you are scored

A name is a maximal run of consecutive words with the same non-O label. A predicted name counts as correct only if its start, end and type all match a true name exactly. Over all sentences and types:

P = \frac{\#\text{correct names}}{\#\text{predicted names}}, \qquad R = \frac{\#\text{correct names}}{\#\text{true names}}, \qquad F_1 = \frac{2PR}{P+R}.

What to aim for. The Evaluation page gives a baseline at 0.66 (ELECTRA fine-tuned on the 50 tagged sentences) and a reference solution at 0.83. On the public leaderboard the "reference" entry has 0.833 and the best contestant 0.812.

First look at the data

The 8 000 count sentences hold 5 387 persons, 5 170 organisations, 5 672 locations, but only 464 creative works and 390 events: two types are rare. Sentences have 1 to 11 names, 2.1 on average.
In the 50 tagged sentences, 117 of the 206 name words are rare (they appear in a single sentence of the whole corpus), and only 22 rare words are not inside a name. Invented words are a strong clue, but 89 name words are ordinary ("city", "court", "festival", "u.s."), so a rare-word rule alone is not enough: "the u.s. food and drug administration" is one organisation of five words, four of them common.
The statement gives hard rules. Names never overlap, two names of the same type are always separated, and a name never starts with "the", "a", "an" or punctuation, and never ends with punctuation or "'s".

A first baseline

Fine-tune ELECTRA on the tagged sentences and take the arg-max label of each word. In our 2-fold check (train on 25 tagged sentences, score on the other 25, then swap), this gives entity $F_1$ of 0.659 and 0.496, mean 0.577. With all 50 sentences, the organisers' baseline reaches 0.66 on the leaderboard. The model sees far too few examples, especially of creative works and events.

The key idea

The counts are a constraint, and a constraint can turn weak labels into strong ones. Take a sentence where the clerk wrote "1 person, 1 location". A model trained on 25 sentences outputs, for every word, a probability for each of the 6 labels. Its arg-max might find two persons, or miss the location. But among all taggings that contain exactly one person and exactly one location and that respect the stated rules, the most probable one is usually right. That gives us a reliable full tagging of the sentence, a pseudo-label, that we can train on.

Finding that best constrained tagging is a small dynamic programme, a Viterbi decoder with an extra counter in its state. Let $\ell_{i}(y)$ be the model's log-probability of label $y$ for word $i$ . We maximise

\sum_{i=1}^{n} \ell_i(y_i)

over label sequences $y_1, \dots, y_n$ . The state after word $i$ is the pair (current label, vector $u$ of names started so far per type). Moving to a type $t$ different from the previous label opens a new name and increments $u_t$ . That move is allowed only if word $i$ may start a name, the previous name may end at word $i-1$ , and $u_t$ stays at most the target count $c_t$ . At the end we require $u = c$ . The number of states is $6 \prod_t (c_t + 1)$ , a few hundred at most, so all 7 500 sentences decode in seconds.

Then we iterate, a classic self-training (EM-like) loop:

train on the gold sentences;
pseudo-label the count sentences with the constrained decoder;
retrain from the pretrained weights on pseudo-labels plus gold (gold repeated 10 times);
go back to 2 with the better model.

Why does this work? The constraint removes most of the model's mistakes before they are trained on. A wrong boundary or a missing name usually breaks the counts, so the decoder is forced to pick another tagging, and the next best is often the correct one. The model then learns from 7 500 mostly correct sentences instead of 25. It is the same reason that a crossword gets easier once you know the length of each word.

At test time there are no counts, so we use the same decoder without the counter. The structural rules still repair impossible outputs, such as a name starting with "the".

The full solution

Freeze the bottom 10 layers of ELECTRA and cache their outputs (float16 on disk, read back with pread). Only the top 2 transformer layers and a linear layer are trained. This was a necessity, not a choice: fine-tuning the whole model used too much memory on our shared machine. It also makes each training round take about 2 minutes on 2 CPU threads.
Round 0: 150 steps on the gold sentences.
Rounds 1 and 2: constrained decoding of the 7 500 count sentences, then retraining (about 1.5 epochs, batch 32, AdamW, learning rate $5\times10^{-5}$ , warm-up then linear decay).
Test: unconstrained structural decoding of the 4 000 sentences.

The heart of the decoder (start of a new name of type t):

valid = used[:, t - 1] >= 1                       # states where one more name of type t fits
src = np.where(valid, np.arange(S) - strides[t - 1], 0)   # state before opening that name
other = ended.copy(); other[t] = NEG              # same type must be separated by a word
sv = other[other[:, src].argmax(0), src] + lp[i, t]

We kept 500 count sentences out of every pseudo-labelling step and used them as a second, larger check: the share of sentences for which the predicted tagging has exactly the right counts.

Results

Entity micro- $F_1$ on the tagged sentences, 2-fold (25 train / 25 scored, then swapped; about 50 names per fold). The official test labels are private.

approach	fold A	fold B	mean	exact counts (500 held out)
gold only (baseline idea)	0.659	0.496	0.577	0.43
1 round, pseudo-labels ignore the counts	0.683			0.48
1 round, count-constrained pseudo-labels	0.775	0.624	0.700	0.56 to 0.59
2 rounds, count-constrained (final recipe)	0.732	0.667	0.699	0.59 to 0.63
final model, all 50 gold sentences, 2 rounds				0.648

Self-training rounds

What the table shows:

The counts are what makes self-training work. Pseudo-labelling with the plain arg-max (red) gains 2.4 points; the same round with count-constrained pseudo-labels gains 11.7 points on the same fold.
A second round helps the count accuracy on both folds (0.56 to 0.59, 0.59 to 0.63) and the $F_1$ of fold B, while fold A's $F_1$ drops. With only about 50 names per fold, one $F_1$ point is half a name, so the gold scores are noisy. The 500-sentence count check, which moves steadily, is the more stable signal.
More gold sentences help. The final model, trained with all 50, gets the counts exactly right on 64.8 % of held-out sentences against 59 to 63 % for the half-gold folds, so its test $F_1$ should be somewhat above our 2-fold estimate of 0.70.

Our estimate, about 0.70 entity $F_1$ , clearly beats the gold-only approach in the same setting (0.58) but stays below the organisers' reference (0.83). Two things we would try next with more compute: fine-tuning more layers (we fine-tuned only the top 2 out of 12 because of the memory limit), and keeping only the pseudo-labels the decoder is confident about (the margin between the best and second-best constrained tagging).

What to remember

Weak labels are often constraints in disguise. Counts, totals, "at least one", bag-level labels: whenever you know a property of the answer, decode the most probable output that satisfies it and train on that.
Viterbi with a counter in the state is the standard tool for "most probable sequence with exactly $k$ of something". The state space is the product of the counters, small here.
Self-training needs a filter. Training on your own arg-max predictions mostly confirms your own mistakes (+2 points); training on constrained predictions corrects them (+12 points).
Validate on more than one signal when the gold set is tiny. About 50 names per fold makes $F_1$ jump by several points at random; a large proxy such as exact-count accuracy is less noisy.
Freeze and cache when compute is short. Caching the frozen bottom of a transformer turns every training round from many minutes into two, which is what makes an iterative method affordable.

Full code (solution.py)

"""AICC Round 11, "Tally": worked solution (weakly supervised NER from mention counts).

1. Fine-tune the top 2 layers of ELECTRA-base (the only allowed model) on the 50 tagged sentences.
   The bottom 10 layers are frozen and their outputs cached once on disk (float16), which keeps
   memory under 3 GB and makes each training round take about two minutes on 2 CPU threads.
2. Pseudo-label the 7 500 count-only sentences with a dynamic programme that finds the most likely
   tagging having EXACTLY the given number of names of each type, and obeying the stated rules
   (no name starts with the/a/an or punctuation, none ends with punctuation or 's).
3. Retrain from the pretrained weights on pseudo-labels + gold, and repeat (2 rounds).
4. Tag the test sentences with the same dynamic programme, without counts.

Validation: 2-fold over the 50 gold sentences (25 to train, 25 to score, then swapped), entity-level
micro F1 as in the competition; plus exact-count accuracy on 500 held-out count sentences.
The official test labels are private.

Usage: python solution.py [A|B|all] [rounds] [1|0]   (run inside the folder holding the csv files)
"""
import os, sys, time, json, copy
import numpy as np
TYPES = ["PERSON", "ORGANIZATION", "LOCATION", "CREATIVE_WORK", "EVENT"]
TAGS = ["O"] + TYPES
NEG = -1e18


def is_punct(w):
    return not any(ch.isalnum() for ch in w)


def decode(words, lp, counts=None):
    """Best label sequence (ints 0..5) under the structural rules, and, if `counts` is given,
    with exactly counts[t] names of each type. lp: n x 6 log-probabilities."""
    n = len(words)
    start_ok = np.array([not (w in ("the", "a", "an") or is_punct(w)) for w in words])
    end_ok = np.array([not (w == "'s" or is_punct(w)) for w in words])
    if counts is None:
        dims = [1] * 5
    else:
        dims = [c + 1 for c in counts]
    strides = np.cumprod([1] + dims[:-1])
    S = int(np.prod(dims))
    used = np.array([[(s // strides[t]) % dims[t] for t in range(5)] for s in range(S)])  # S x 5
    score = np.full((6, S), NEG)
    back = np.zeros((n, 6, S, 2), dtype=np.int64)       # (prev label, prev state)
    # position 0
    score[0, 0] = lp[0, 0]
    if start_ok[0]:
        for t in range(1, 6):
            s = 0 if counts is None else strides[t - 1]
            if counts is None or dims[t - 1] > 1:
                score[t, s] = lp[0, t]
    for i in range(1, n):
        new = np.full((6, S), NEG)
        # candidates "leave previous label": prev label p may end at i-1 only if end_ok
        ended = score.copy()
        if not end_ok[i - 1]:
            ended[1:] = NEG
        # label O
        best_p = ended.argmax(0)
        new[0] = ended[best_p, np.arange(S)] + lp[i, 0]
        back[i, 0, :, 0] = best_p
        back[i, 0, :, 1] = np.arange(S)
        for t in range(1, 6):
            # continuation
            cont = score[t] + lp[i, t]
            cand_val, cand_p, cand_s = cont, np.full(S, t), np.arange(S)
            if start_ok[i]:
                other = ended.copy()
                other[t] = NEG                                  # same type must be separated
                if counts is None:
                    src = np.arange(S)
                    valid = np.ones(S, bool)
                else:
                    valid = used[:, t - 1] >= 1
                    src = np.where(valid, np.arange(S) - strides[t - 1], 0)
                bp = other[:, src].argmax(0)
                sv = other[bp, src] + lp[i, t]
                sv[~valid] = NEG
                better = sv > cand_val
                cand_val = np.where(better, sv, cand_val)
                cand_p = np.where(better, bp, cand_p)
                cand_s = np.where(better, src, cand_s)
            new[t] = cand_val
            back[i, t, :, 0] = cand_p
            back[i, t, :, 1] = cand_s
        score = new
    final = score.copy()
    if not end_ok[n - 1]:
        final[1:] = NEG
    s_end = S - 1 if counts is not None else 0
    lab = int(final[:, s_end].argmax())
    if final[lab, s_end] <= NEG / 2:
        return None                                              # infeasible
    out = [lab]
    s = s_end
    for i in range(n - 1, 0, -1):
        p, s = back[i, lab, s]
        lab = int(p)
        out.append(lab)
    return out[::-1]


def spans(tags):
    """Set of (start, end, type) for a list of int labels."""
    out, i = set(), 0
    while i < len(tags):
        if tags[i] != 0:
            j = i
            while j + 1 < len(tags) and tags[j + 1] == tags[i]:
                j += 1
            out.add((i, j, tags[i]))
            i = j + 1
        else:
            i += 1
    return out


def f1(pred_list, gold_list):
    tp = npred = ngold = 0
    for p, g in zip(pred_list, gold_list):
        P, G = spans(p), spans(g)
        tp += len(P & G); npred += len(P); ngold += len(G)
    pr = tp / max(npred, 1); rc = tp / max(ngold, 1)
    return 2 * pr * rc / max(pr + rc, 1e-9)




os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("HF_HOME", os.path.join(os.getcwd(), "hf"))
import numpy as np, pandas as pd, torch
import torch.nn as nn, torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
torch.set_num_threads(2)
FOLD = sys.argv[1] if len(sys.argv) > 1 else "all"           # A / B: 2-fold check on the 50 gold sentences
ROUNDS = int(sys.argv[2]) if len(sys.argv) > 2 else 2         # self-training rounds
CONSTRAINED = (sys.argv[3] if len(sys.argv) > 3 else "1") == "1"   # use the counts when pseudo-labelling
TOPK = 2                                  # number of top ELECTRA layers we fine-tune
NAME = "google/electra-base-discriminator"
tok = AutoTokenizer.from_pretrained(NAME)
from transformers import AutoConfig
from transformers.models.electra.modeling_electra import ElectraLayer
from safetensors.torch import save_file, load_file
CFG = AutoConfig.from_pretrained(NAME)
rng = np.random.default_rng(0); torch.manual_seed(0)

gold = pd.read_csv("train_tags.csv"); tc = pd.read_csv("train_counts.csv"); te = pd.read_csv("test.csv")
gw = gold.text.str.split().tolist(); gy = [[TAGS.index(t) for t in s.split()] for s in gold.tags]
cw = tc.text.str.split().tolist(); cc = tc[TYPES].values.tolist(); tw = te.text.str.split().tolist()
ALL = gw + cw + tw                          # sentence k of the cache
G0, C0, T0 = 0, len(gw), len(gw) + len(cw)
log = []
def say(*a):
    print(*a, flush=True); log.append(" ".join(map(str, a)))

# ------------------------------------------------ cache the output of the frozen bottom layers
CACHE = f"cache_l{12-TOPK}.f16"
if not os.path.exists(CACHE + ".json"):
    base = AutoModel.from_pretrained(NAME).eval()
    save_file({k: v.contiguous() for k, v in base.encoder.layer[12 - TOPK:].state_dict().items()}, "top_layers.safetensors")
    offs, firsts, chunks, pos = [], [], [], 0
    mm = open(CACHE, "wb")
    with torch.no_grad():
        for i in range(0, len(ALL), 32):
            b = ALL[i:i + 32]
            enc = tok(b, is_split_into_words=True, padding=True, truncation=True, max_length=128, return_tensors="pt")
            hs = base(**enc, output_hidden_states=True).hidden_states[12 - TOPK]
            for k in range(len(b)):
                n = int(enc["attention_mask"][k].sum()); wid = enc.word_ids(k); fst = {}
                for s, w in enumerate(wid):
                    if w is not None and w not in fst: fst[w] = s
                firsts.append([fst.get(w, 0) for w in range(len(b[k]))])
                mm.write(hs[k, :n].numpy().astype(np.float16).tobytes()); offs.append((pos, n)); pos += n
            if i % 3200 == 0: print("cache", i, flush=True)
    mm.close(); json.dump({"offs": offs, "firsts": firsts, "total": pos}, open(CACHE + ".json", "w"))
if "base" in globals(): del base
TEMPLATE = nn.ModuleList([ElectraLayer(CFG) for _ in range(TOPK)]); TEMPLATE.load_state_dict(load_file("top_layers.safetensors")); TEMPLATE.eval()
meta = json.load(open(CACHE + ".json")); OFFS, FIRSTS = meta["offs"], meta["firsts"]
FD = os.open(CACHE, os.O_RDONLY)          # read rows with pread: no memory-mapped pages piling up
def rows(p, m):
    return np.frombuffer(os.pread(FD, m * 768 * 2, p * 768 * 2), dtype=np.float16).reshape(m, 768)


class Head(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = copy.deepcopy(TEMPLATE); self.drop = nn.Dropout(0.1); self.out = nn.Linear(768, 6)
    def forward(self, h, mask):
        ext = (1.0 - mask[:, None, None, :].float()) * torch.finfo(torch.float32).min
        for L in self.layers:
            o = L(h, attention_mask=ext); h = o[0] if isinstance(o, tuple) else o
        return self.out(self.drop(h))


def batch(ids):
    n = max(OFFS[k][1] for k in ids); n = 32 * ((n + 31) // 32)          # fixed bucket sizes keep the allocator from fragmenting
    h = torch.zeros(len(ids), n, 768); mask = torch.zeros(len(ids), n)
    for b, k in enumerate(ids):
        p, m = OFFS[k]; h[b, :m] = torch.from_numpy(rows(p, m).astype(np.float32)); mask[b, :m] = 1
    return h, mask


def train(ids, labels, epochs=1.5, bs=32, lr=5e-5):
    model = Head(); opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    steps = max(150, int(epochs * len(ids) / bs))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1, (s + 1) / 30) * max(0.0, 1 - s / steps))
    model.train(); order = []
    for step in range(steps):
        if len(order) < bs: order += list(rng.permutation(len(ids)))
        sel, order = order[:bs], order[bs:]
        h, mask = batch([ids[j] for j in sel]); lab = torch.full(mask.shape, -100, dtype=torch.long)
        for b, j in enumerate(sel):
            lab[b, FIRSTS[ids[j]]] = torch.tensor(labels[j])
        loss = F.cross_entropy(model(h, mask).flatten(0, 1), lab.flatten(), ignore_index=-100)
        opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    say(f"  trained {steps} steps, last loss {loss.item():.3f}")
    return model.eval()


@torch.no_grad()
def logprobs(model, ids, bs=16):
    out = []
    for i in range(0, len(ids), bs):
        b = ids[i:i + bs]; h, mask = batch(b); lp = F.log_softmax(model(h, mask), -1)
        out += [lp[k, FIRSTS[j]].numpy() for k, j in enumerate(b)]
    return out


perm = rng.permutation(len(gw))
gtr, gev = {"A": (perm[:25], perm[25:]), "B": (perm[25:], perm[:25]), "all": (perm, [])}[FOLD]
cperm = rng.permutation(len(cw)); cev, ctr = cperm[:500], cperm[500:]


def evaluate(model, tag):
    res = {}
    if len(gev):
        lp = logprobs(model, [G0 + i for i in gev])
        res["gold_f1"] = round(f1([decode(gw[i], l) for i, l in zip(gev, lp)], [gy[i] for i in gev]), 4)
        res["gold_f1_argmax"] = round(f1([list(l.argmax(1)) for l in lp], [gy[i] for i in gev]), 4)
    lp = logprobs(model, [C0 + i for i in cev])
    cnt = [[sum(1 for s in spans(decode(cw[i], l)) if s[2] == t) for t in range(1, 6)] for i, l in zip(cev, lp)]
    res["count_exact"] = round(float(np.mean([c == cc[i] for c, i in zip(cnt, cev)])), 4)
    say(tag, json.dumps(res)); return res


if __name__ == "__main__":
    t0 = time.time(); say(f"fold {FOLD}, {len(gtr)} gold train, constrained={CONSTRAINED}")
    model = train([G0 + i for i in gtr], [gy[i] for i in gtr], epochs=150)
    R = {"round0": evaluate(model, "round 0 (gold only)")}
    for r in range(1, ROUNDS + 1):
        lp = logprobs(model, [C0 + i for i in ctr]); ids, labs = [], []
        for i, l in zip(ctr, lp):
            d = decode(cw[i], l, cc[i] if CONSTRAINED else None)
            if d is not None: ids.append(C0 + i); labs.append(d)
        del lp
        ids += [G0 + i for i in gtr] * 10; labs += [gy[i] for i in gtr] * 10
        say(f"round {r}: {len(ids)} training sentences, {time.time()-t0:.0f}s")
        model = train(ids, labs)
        R[f"round{r}"] = evaluate(model, f"round {r}")
    say(f"total {time.time()-t0:.0f}s")
    json.dump({"results": R, "log": log}, open(f"res2_{FOLD}_{int(CONSTRAINED)}.json", "w"), indent=1)
    if FOLD == "all":
        lp = logprobs(model, list(range(T0, T0 + len(tw))))
        pred = [decode(w, l) for w, l in zip(tw, lp)]
        pd.DataFrame({"id": te.id, "tags": [" ".join(TAGS[t] for t in p) for p in pred]}).to_csv("submission.csv", index=False)