← Back to the problem

AICC (IOAI Community Contest) 2026 · Solution

Word Lookups

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

Formatted by IOAI Training from the original (AI-Community-Contest/solutions, licence MIT). The original is authoritative.

Loading and preprocessing

import ast
import pandas as pd
import torch
from collections import Counter
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence

# -------------------------
# Load CSVs
# -------------------------
train_df = pd.read_csv("train.csv")
test_df = pd.read_csv("test.csv")

# Convert string representations back to lists
train_df["chars"] = train_df["chars"].apply(ast.literal_eval)
test_df["chars"] = test_df["chars"].apply(ast.literal_eval)

train_chars = train_df["chars"].tolist()
test_chars = test_df["chars"].tolist()

# -------------------------
# Build vocabulary from TRAIN only
# -------------------------
all_chars = train_chars + test_chars
counter = Counter(c for sent in all_chars for c in sent)

char2id = {"PAD": 0, "UNK": 1}
for i, c in enumerate(counter.keys(), start=2):
    char2id[c] = i

vocab_size = len(char2id)

# -------------------------
# Encode
# -------------------------
def encode(sent):
    return torch.tensor([char2id.get(c, 1) for c in sent])

train_data = [encode(s) for s in train_chars]
test_data = [encode(s) for s in test_chars]

# -------------------------
# Dataset
# -------------------------
class CharDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

PAD_ID = 0

def collate_fn(batch):
    return pad_sequence(batch, batch_first=True, padding_value=PAD_ID)

# -------------------------
# Datasets
# -------------------------
train_dataset = CharDataset(train_data)
test_dataset = CharDataset(test_data)

# -------------------------
# DataLoaders
# -------------------------
train_loader = DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
    collate_fn=collate_fn,
)

test_loader = DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=False,
    collate_fn=collate_fn,
)
from torch.utils.data import ConcatDataset, DataLoader

combined_dataset = ConcatDataset([
    train_dataset,
    test_dataset
])

combined_loader = DataLoader(
    combined_dataset,
    batch_size=16,
    shuffle=True,
    collate_fn=collate_fn
)

BiLSTM training

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
import torch.nn as nn
import torch.nn.functional as F

class BiLM(nn.Module):
    def __init__(self, vocab_size, emb_dim=128, hidden_dim=256):
        super().__init__()

        self.embedding = nn.Embedding(
            vocab_size,
            emb_dim,
            padding_idx=PAD_ID
        )

        self.forward_lstm = nn.LSTM(
            emb_dim,
            hidden_dim,
            batch_first=True
        )

        self.backward_lstm = nn.LSTM(
            emb_dim,
            hidden_dim,
            batch_first=True
        )

        self.forward_out = nn.Linear(hidden_dim, vocab_size)
        self.backward_out = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):

        emb = self.embedding(x)

        fwd_hidden, _ = self.forward_lstm(emb)

        rev_emb = torch.flip(emb, dims=[1])
        bwd_hidden, _ = self.backward_lstm(rev_emb)

        fwd_logits = self.forward_out(fwd_hidden)
        bwd_logits = self.backward_out(bwd_hidden)

        return fwd_logits, bwd_logits
def bilm_loss(fwd_logits, bwd_logits, x):

    # ---------- Forward LM ----------

    fwd_target = x[:, 1:]
    fwd_pred = fwd_logits[:, :-1]

    loss_fwd = F.cross_entropy(
        fwd_pred.reshape(-1, fwd_pred.size(-1)),
        fwd_target.reshape(-1),
        ignore_index=PAD_ID
    )

    # ---------- Backward LM ----------

    x_rev = torch.flip(x, dims=[1])

    bwd_target = x_rev[:, 1:]
    bwd_pred = bwd_logits[:, :-1]

    loss_bwd = F.cross_entropy(
        bwd_pred.reshape(-1, bwd_pred.size(-1)),
        bwd_target.reshape(-1),
        ignore_index=PAD_ID
    )

    return loss_fwd + loss_bwd
import os
from tqdm.notebook import tqdm

model_path = "combmodel (1).pth"
model = BiLM(vocab_size).to(device)

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3
)

EPOCHS = 0 if os.path.exists(model_path) else 10
if os.path.exists(model_path):
    model.load_state_dict(torch.load(model_path, map_location=device))
    print(f"Loaded {model_path}")

for epoch in tqdm(range(EPOCHS)):

    model.train()

    running_loss = 0.0

    for batch in combined_loader:

        batch = batch.to(device)

        optimizer.zero_grad()

        fwd_logits, bwd_logits = model(batch)

        loss = bilm_loss(
            fwd_logits,
            bwd_logits,
            batch
        )

        loss.backward()

        torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            5.0
        )

        optimizer.step()

        running_loss += loss.item()

    print(
        f"Epoch {epoch+1}: "
        f"{running_loss/len(combined_loader):.4f}"
    )
Loaded combmodel (1).pth
0it [00:00, ?it/s]
torch.save(model.state_dict(), model_path)

Boundaries

model.eval()
BiLM(
  (embedding): Embedding(4397, 128, padding_idx=0)
  (forward_lstm): LSTM(128, 256, batch_first=True)
  (backward_lstm): LSTM(128, 256, batch_first=True)
  (forward_out): Linear(in_features=256, out_features=4397, bias=True)
  (backward_out): Linear(in_features=256, out_features=4397, bias=True)
)
import torch
import torch.nn.functional as F

@torch.no_grad()
def get_lm_outputs(model, sentence):
    """
    sentence: 1D tensor of character ids (length T)

    Returns:
        fwd_probs: (T, V)
        bwd_probs: (T, V)
    """

    device = next(model.parameters()).device

    x = sentence.unsqueeze(0).to(device)

    fwd_logits, bwd_logits = model(x)

    fwd_probs = F.softmax(fwd_logits[0], dim=-1)
    bwd_probs = F.softmax(bwd_logits[0], dim=-1)

    return fwd_probs.cpu(), bwd_probs.cpu()
def compute_surprisal(sentence):

    fwd_probs, bwd_probs = get_lm_outputs(model, sentence)

    x = sentence.cpu()

    T = len(x)

    fwd = torch.zeros(T - 1)
    bwd = torch.zeros(T - 1)

    # forward LM
    for i in range(T - 1):
        target = x[i + 1]
        p = fwd_probs[i, target]
        fwd[i] = -torch.log(p + 1e-12)

    # backward LM
    x_rev = torch.flip(x, dims=[0])

    for i in range(T - 1):
        target = x_rev[i + 1]
        p = bwd_probs[i, target]
        bwd[T - 2 - i] = -torch.log(p + 1e-12)

    return fwd, bwd
def boundary_scores(sentence):

    fwd, bwd = compute_surprisal(sentence)

    return fwd + bwd
scores = boundary_scores(train_data[0])

print(scores)
tensor([ 2.4260,  7.6363,  7.8438,  3.7121, 13.9421,  4.9846,  9.6920,  6.9366,
         7.5610,  6.6622,  5.6499,  4.5213,  7.0375,  2.3595, 12.2140,  1.0803,
         7.2117,  6.9750,  0.7285, 19.3950,  9.7396, 12.3365,  8.4980,  6.8240,
         9.0279,  9.8234, 12.0777, 10.6452, 14.4207])
import torch

all_scores = []

for i in tqdm(range(len(train_data))):
    scores = boundary_scores(train_data[i])
    all_scores.append(scores)

all_scores = torch.cat(all_scores)

mean_score = all_scores.mean()

print(mean_score)
  0%|          | 0/37086 [00:00<?, ?it/s]tensor(7.1728)
scores = all_scores

Some of the following cells require access to solution.csv

sol_df = pd.read_csv("solution.csv")
bio_tags = ast.literal_eval(sol_df.iloc[0]["bio_tags"])
print(bio_tags)
['S', 'B', 'E', 'B', 'E', 'B', 'E', 'B', 'E', 'S', 'B', 'E', 'B', 'E', 'S', 'B', 'E', 'B', 'E', 'S', 'B', 'E', 'B', 'E', 'B', 'E', 'B', 'E', 'S', 'B', 'E', 'S', 'B', 'E', 'B', 'E', 'S', 'B', 'E', 'B', 'E', 'S']
id2char = {v: k for k, v in char2id.items()}
idx = 2
Y = ast.literal_eval(sol_df.iloc[idx]["bio_tags"])
X = test_data

print(X[idx])
print(Y)
tensor([1588,  971,   34, 1130, 1131,    4,  460, 3042, 3042, 4142, 4142, 1807,
        1325,   29,  234,   34,  171,   35,   26])
['B', 'E', 'S', 'B', 'E', 'S', 'S', 'B', 'M', 'M', 'E', 'B', 'E', 'S', 'S', 'S', 'B', 'E', 'S']
def visualize_with_gold(idx):
    Y = ast.literal_eval(sol_df.iloc[idx]["bio_tags"])
    X = test_data

    chars = X[idx]
    tags = Y

    scores = boundary_scores(X[idx])

    print("Sentence:\n")

    for i, ch in enumerate(chars):
        print(id2char[ch.item()], end="")

        # print score after character i
        if i < len(scores):
            print(f" [{scores[i]:.2f}] ", end="")

    print("\n")

    print("Gold:\n")

    for ch, tag in zip(chars, tags):
        print(id2char[ch.item()], end="")

        # A word ends after S or E
        if tag in ("S", "E"):
            print(" | ", end="")

    print()
visualize_with_gold(0)
Sentence:

据 [4.05] 此 [0.98] 间 [8.56] 新 [0.21] 闻 [0.67] 媒 [0.20] 体 [4.50] 报 [0.21] 道 [4.22] , [5.52] 此 [2.68] 次 [13.28] 视 [4.13] 察 [10.30] 的 [4.89] 主 [1.47] 要 [7.13] 目 [2.06] 的 [6.11] 是 [13.56] 了 [3.19] 解 [17.77] 位 [6.82] 于 [10.82] 冲 [1.14] 突 [15.34] 地 [2.36] 区 [12.34] 及 [9.59] 近 [15.06] 处 [9.60] 的 [9.28] 边 [1.52] 防 [4.07] 部 [2.15] 队 [9.26] 的 [5.16] 战 [16.65] 备 [11.89] 情 [0.24] 况 [7.17] 。

Gold:

据 | 此间 | 新闻 | 媒体 | 报道 | , | 此次 | 视察 | 的 | 主要 | 目的 | 是 | 了解 | 位于 | 冲突 | 地区 | 及 | 近处 | 的 | 边防 | 部队 | 的 | 战备 | 情况 | 。 |

Decoding

def decode_threshold(sentence, threshold=None):
    scores = boundary_scores(sentence)

    if threshold is None:
        threshold = scores.mean() + scores.std()

    chars = [id2char[i] for i in sentence.tolist()]

    words = []
    current = []

    for i, ch in enumerate(chars):
        current.append(ch)

        if i < len(scores) and scores[i] >= threshold:
            words.append("".join(current))
            current = []

    if current:
        words.append("".join(current))

    return words
print(decode_threshold(test_data[0]))
#print(df.iloc[0]["words"])
['据此间新闻媒体报道,此次', '视察的主要目的是', '了解', '位于冲突', '地区', '及近', '处的边防部队的战', '备', '情况。']
threshold = scores.mean()
print(decode_threshold(test_data[4], threshold=threshold))
#print(df.iloc[0]["words"])
['针对此案', '圆满', '解决', '一', '事,', '秘书长说:“', '对', '老虎', '吃', '牛', '这个', '特殊', '行政', '补偿', '案件', '的', '妥善处理,将', '对', '野生', '动物保护', '事业', '带来', '积极', '影响', '。”']
def boundaries_from_words(words):
    b = set()
    pos = 0
    for w in words[:-1]:
        pos += len(w)
        b.add(pos)
    return b

def boundary_f1(pred_words, gold_words):
    pred = boundaries_from_words(pred_words)
    gold = boundaries_from_words(gold_words)

    tp = len(pred & gold)

    precision = tp / len(pred) if pred else 0
    recall = tp / len(gold) if gold else 0

    if precision + recall == 0:
        f1 = 0
    else:
        f1 = 2 * precision * recall / (precision + recall)

    return precision, recall, f1
i = 14
threshold = 2 # scores.mean() - 2 * scores.std()

pred = decode_threshold(test_data[i], threshold=threshold)
gold = ast.literal_eval(sol_df.iloc[i]["bio_tags"])


p, r, f1 = boundary_f1(pred, gold)

print(pred)
print(gold)
print(f"P={p:.3f} R={r:.3f} F1={f1:.3f}")
['短', '短', '中', ',', '她', '竟', '在', '比赛', '中', '跑', '进', ',', '可', '谓', '前', '无', '古', '人', '。']
['B', 'E', 'S', 'S', 'S', 'S', 'S', 'B', 'E', 'S', 'S', 'S', 'S', 'B', 'E', 'B', 'M', 'M', 'E', 'S']
P=1.000 R=0.947 F1=0.973
threshold =  scores.mean()# - 3 * scores.std()
total_tp = 0
total_pred = 0
total_gold = 0

for i in tqdm(range(len(test_data))):
    pred_words = decode_threshold(test_data[i], threshold=threshold)
    gold_words = ast.literal_eval(sol_df.iloc[i]["bio_tags"])

    pred = boundaries_from_words(pred_words)
    gold = boundaries_from_words(gold_words)

    total_tp += len(pred & gold)
    total_pred += len(pred)
    total_gold += len(gold)

precision = total_tp / total_pred if total_pred else 0
recall = total_tp / total_gold if total_gold else 0

if precision + recall == 0:
    f1 = 0
else:
    f1 = 2 * precision * recall / (precision + recall)

print("Precision:", precision)
print("Recall:", recall)
print("F1:", f1)
  0%|          | 0/9272 [00:00<?, ?it/s]
import pandas as pd

def words_to_bmes(words):
    tags = []

    for word in words:
        if len(word) == 1:
            tags.append("S")
        else:
            tags.append("B")
            tags.extend(["M"] * (len(word) - 2))
            tags.append("E")

    return tags


pred_tags = []

for chars in test_data:
    # your threshold segmentatio
    words = decode_threshold(chars, threshold=scores.mean() - 0.5 * scores.std())

    # convert words -> BMES tags
    tags = words_to_bmes(words)

    pred_tags.append(str(tags))


submission = pd.DataFrame({
    "id": test_df["id"],
    "bio_tags": pred_tags
})

submission.to_csv("submission.csv", index=False)

submission.head()
      id                                           bio_tags
0  10259  ['B', 'M', 'E', 'B', 'M', 'M', 'M', 'M', 'M', ...
1  10962  ['B', 'E', 'B', 'E', 'B', 'E', 'S', 'B', 'E', ...
2    884  ['S', 'S', 'S', 'B', 'E', 'S', 'S', 'B', 'M', ...
3  17894  ['S', 'S', 'S', 'S', 'S', 'S', 'B', 'E', 'B', ...
4  38013  ['B', 'E', 'S', 'S', 'B', 'E', 'B', 'E', 'S', ...