Checklist OAI 2024 Stage I · Task 3
SOTA-AI Editorial
Dependency Parsing
Build dependency trees for Polish sentences from HerBERT word vectors by training small models that predict tree distances and depths.
Hints
Hint 1
Parse the CoNLL head column and check how many roots each sentence has.
Hint 2
Dependency structure can be represented through pairwise tree distances and each token's depth from the root.
Hint 3
Contextual word representations can provide features for both distance and root-depth prediction.
Hint 4
Recover an undirected tree from predicted pairwise distances with an MST, then choose the token with minimum predicted depth as root.
The solution in one line
Use frozen HerBERT word vectors to predict pairwise tree distances and root depths, then recover the tree with an MST and choose the shallowest token as root.
The editorial
Dependency Parsing — Stage I (2024)
Problem name: Dependency Parsing (dependency_parsing)
Domain: Polish NLP, unsupervised-style dependency induction with labeled training data
Metric: Average UUAS (unlabeled undirected attachment score) and root-placement accuracy, each mapped linearly from 0.50 to 0.85 and clipped at that range. The sum is worth up to 2 points; full credit requires both metrics to be at least 0.85.
Abridged statement
Train a parser for Polish sentences using the released 1,000-sentence labeled training set. Predict a dependency tree and root position for each sentence. The official interface permits the listed HerBERT pretrained model and requires CPU evaluation; training and evaluation have organizer time limits. The 200-sentence validation split is used only for evaluation.
Dataset analysis and EDA
The official repository contains train.conll (1,000 sentences) and valid.conll (200 sentences), plus the organizer's utils.py and validation_script.py. These assets were retrieved for inspection and not copied into the solution directory. The train set has 9,663 tokens (mean 9.66 words/sentence, median 8, maximum 39); validation has 1,821 tokens (mean 9.11, median 8, maximum 45). The annotated root is among the first three positions in 53.4% of training sentences and 55.0% of validation sentences. The code reads gold heads from the CoNLL files and constructs undirected tree-distance and root-depth targets from training sentences only.
Experiments
- Official starter / evaluator inspection: reviewed the starter notebook, its TODO interface, the supplied split files, and
validation_script.py. No official worked solution or organizer-reported score was found in the inspected materials. - Data-only baselines: no heuristic baseline was assigned a validation score. Root-position frequencies show only a weak early-root prior; they are descriptive, not a measured parser result.
- HerBERT representation plus distance/depth regressors: implemented as the selected route. The frozen
allegro/herbert-base-casedword representations feed one symmetric pairwise tree-distance regressor and one root-depth regressor. A minimum spanning tree recovers undirected edges and the predicted minimum-depth token is selected as root. This code was not run because neithertransformersnor the allowed HerBERT checkpoint/tokenizer is installed or cached locally, and shell network access is unavailable.
Selected solution
The candidate implementation is dependency_parsing.py. It loads the official HerBERT checkpoint through Transformers, averages wordpiece vectors for each word, fits the distance and depth heads from train.conll, saves compact head state dictionaries, and evaluates against valid.conll. It does not load validation data during training. The source is a complete execution route conditional on the organizer-allowed Transformers package and pretrained checkpoint being available; it has not been dependency-resolved or runtime-verified in this environment.
Validation score and evidence
No local validation score was measured. The exact blocker is that this environment has no transformers installation and no cached allegro/herbert-base-cased weights/tokenizer, while outbound shell downloads are unavailable. The official 200-sentence validation split and scoring script are available, but scoring without the required model would not be meaningful. No organizer-reported score was identified. Therefore there is no basis to claim full credit or hidden-test performance.
Compute and runtime
The supplied route freezes HerBERT and trains only two small regression heads, with CPU selected automatically when CUDA is unavailable. The code defaults to 40 epochs per head. Because the pretrained model could not be loaded, training time, CPU evaluation time, and compliance with the official training/evaluation time limits remain unverified. The task allows only CPU evaluation; no GPU-specific route is required.
Alternatives
Potential routes include a majority root-position prior, a supervised biaffine arc parser with HerBERT features, or a transition-based parser. None was trained or scored here. A cached organizer-allowed HerBERT checkpoint is required to evaluate the selected approach and make a fair comparison.
Code
Solution code
"""HerBERT-based distance/depth parser for Polish AI Olympiad 2024, Stage I.
This standalone implementation needs ``torch``, ``transformers`` and NumPy,
plus the organizer's ``train.conll`` / ``valid.conll`` files. It only trains the
small task heads on ``train.conll``; HerBERT is frozen. The official HerBERT
checkpoint is downloaded by ``from_pretrained`` when it is not already cached.
"""
from __future__ import annotations
import argparse
from pathlib import Path
from typing import Iterable, Sequence
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from transformers import AutoModel, AutoTokenizer
HERBERT_MODEL = "allegro/herbert-base-cased"
DISTANCE_MODEL_PATH = "distance_model.pth"
DEPTH_MODEL_PATH = "depth_model.pth"
class Sentence:
def __init__(self, words: Sequence[str]):
self.words = list(words)
def __len__(self) -> int:
return len(self.words)
def __repr__(self) -> str:
return " ".join(self.words)
class ParsedSentence(Sentence):
"""Organizer-compatible sentence representation with 1-based heads."""
def __init__(self, words: Sequence[str], heads: Sequence[int]):
if len(words) != len(heads):
raise ValueError("Different number of words and heads")
if not all(0 <= head <= len(words) for head in heads):
raise ValueError("Invalid head index")
if list(heads).count(0) != 1:
raise ValueError("A dependency tree must have exactly one root")
super().__init__(words)
self.heads = list(map(int, heads))
self.root = self.heads.index(0)
@classmethod
def from_edges_and_root(
cls, words: Sequence[str], edges: Iterable[tuple[int, int]], root: int
) -> "ParsedSentence":
n = len(words)
if not 0 <= root < n:
raise ValueError("Root index is outside the sentence")
adjacency = [[] for _ in range(n)]
for i, j in edges:
if i == j or not (0 <= i < n and 0 <= j < n):
raise ValueError("Invalid tree edge")
adjacency[i].append(j)
adjacency[j].append(i)
heads = [-1] * n
heads[root] = 0
stack = [(root, -1)]
while stack:
node, parent = stack.pop()
for child in adjacency[node]:
if child == parent:
continue
if heads[child] != -1:
raise ValueError("Edges contain a cycle")
heads[child] = node + 1
stack.append((child, node))
if any(head == -1 for head in heads):
raise ValueError("Edges are not connected")
return cls(words, heads)
def get_sorted_edges(self) -> list[tuple[int, int]]:
return [
(min(i, head - 1), max(i, head - 1))
for i, head in enumerate(self.heads)
if head != 0
]
def read_conll(path: str | Path) -> list[ParsedSentence]:
sentences: list[ParsedSentence] = []
words: list[str] = []
heads: list[int] = []
def flush() -> None:
nonlocal words, heads
if words:
sentences.append(ParsedSentence(words, heads))
words, heads = [], []
with Path(path).open(encoding="utf-8") as source:
for line in source:
if not line.strip():
flush()
continue
if line.startswith("#"):
continue
fields = line.rstrip("\n").split("\t")
if len(fields) < 7:
continue
try:
int(fields[0])
head = int(fields[6])
except ValueError:
# Skip CoNLL range/empty-node rows, which are not tokens.
continue
words.append(fields[1])
heads.append(head)
flush()
return sentences
def uuas_score(golden: ParsedSentence, predicted: ParsedSentence) -> float:
gold_edges = set(golden.get_sorted_edges())
pred_edges = set(predicted.get_sorted_edges())
return len(gold_edges & pred_edges) / max(1, len(gold_edges))
def get_distances(sentence: ParsedSentence) -> np.ndarray:
"""Return the shortest-path distance between every pair of tree words."""
n = len(sentence)
adjacency = [[] for _ in range(n)]
for child, head in enumerate(sentence.heads):
if head:
parent = head - 1
adjacency[child].append(parent)
adjacency[parent].append(child)
distances = np.zeros((n, n), dtype=np.float32)
for source in range(n):
distances[source, :] = np.inf
distances[source, source] = 0.0
queue = [source]
for node in queue:
for child in adjacency[node]:
if not np.isfinite(distances[source, child]):
distances[source, child] = distances[source, node] + 1.0
queue.append(child)
return distances
def get_word_embeddings(
sentences: Sequence[Sentence], tokenizer, model: nn.Module, batch_size: int = 24
) -> list[torch.Tensor]:
"""Extract frozen contextual HerBERT vectors and average word subwords."""
model.eval()
device = next(model.parameters()).device
results: list[torch.Tensor] = []
texts = [" ".join(sentence.words) for sentence in sentences]
with torch.no_grad():
for start in range(0, len(sentences), batch_size):
batch_sentences = sentences[start : start + batch_size]
encoded = tokenizer(
texts[start : start + batch_size],
add_special_tokens=False,
padding=True,
truncation=True,
return_tensors="pt",
)
encoded = {key: value.to(device) for key, value in encoded.items()}
hidden = model(**encoded).last_hidden_state
for row, sentence in enumerate(batch_sentences):
length = int(encoded["attention_mask"][row].sum().item())
sentence_ids = encoded["input_ids"][row, :length].tolist()
word_ids: list[list[int]] = [
tokenizer(word, add_special_tokens=False)["input_ids"]
for word in sentence.words
]
joined_ids = [token_id for group in word_ids for token_id in group]
if joined_ids != sentence_ids:
raise ValueError(
"HerBERT word tokenization did not match sentence tokenization; "
f"sentence={sentence!r}"
)
vectors = []
cursor = 0
for ids in word_ids:
end = cursor + len(ids)
if end == cursor:
vectors.append(hidden.new_zeros(hidden.shape[-1]))
else:
vectors.append(hidden[row, cursor:end].mean(dim=0))
cursor = end
results.append(torch.stack(vectors).cpu())
return results
class DistanceModel(nn.Module):
"""Symmetric regression head for pairwise tree distances."""
def __init__(self, embedding_dim: int = 768, hidden_dim: int = 64):
super().__init__()
self.projection = nn.Sequential(nn.Linear(embedding_dim, hidden_dim), nn.GELU())
self.pair_head = nn.Sequential(
nn.Linear(hidden_dim * 2, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, 1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
projected = self.projection(x)
left = projected.unsqueeze(2)
right = projected.unsqueeze(1)
pair = torch.cat((torch.abs(left - right), left * right), dim=-1)
distances = F.softplus(self.pair_head(pair).squeeze(-1))
length = distances.shape[1]
return distances.masked_fill(torch.eye(length, device=x.device, dtype=torch.bool)[None], 0.0)
class DepthModel(nn.Module):
"""Predict each word's distance from the syntactic root."""
def __init__(self, embedding_dim: int = 768, hidden_dim: int = 128):
super().__init__()
self.network = nn.Sequential(
nn.Linear(embedding_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.GELU(),
nn.Linear(hidden_dim // 2, 1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.softplus(self.network(x))
def _collate(examples, target_kind: str):
embeddings, targets = zip(*examples)
batch = len(examples)
max_len = max(len(item) for item in embeddings)
dim = embeddings[0].shape[-1]
x = torch.zeros(batch, max_len, dim, dtype=torch.float32)
shape = (batch, max_len, max_len) if target_kind == "distance" else (batch, max_len, 1)
y = torch.full(shape, float("inf"), dtype=torch.float32)
for i, (embedding, target) in enumerate(zip(embeddings, targets)):
length = len(embedding)
x[i, :length] = embedding
if target_kind == "distance":
y[i, :length, :length] = torch.as_tensor(target, dtype=torch.float32)
else:
y[i, :length, :] = torch.as_tensor(target, dtype=torch.float32)
mask = torch.isfinite(y)
return x, y, mask
def _fit_regressor(
model: nn.Module,
embeddings: Sequence[torch.Tensor],
targets: Sequence[np.ndarray],
target_kind: str,
epochs: int,
learning_rate: float,
batch_size: int = 32,
) -> nn.Module:
device = next(model.parameters()).device
order = np.arange(len(embeddings))
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4)
model.train()
for epoch in range(epochs):
np.random.shuffle(order)
epoch_loss = 0.0
batches = 0
for start in range(0, len(order), batch_size):
ids = order[start : start + batch_size]
x, y, mask = _collate(
[(embeddings[i], targets[i]) for i in ids], target_kind
)
x, y, mask = x.to(device), y.to(device), mask.to(device)
optimizer.zero_grad(set_to_none=True)
prediction = model(x)
loss = ((prediction[mask] - y[mask]) ** 2).mean()
loss.backward()
optimizer.step()
epoch_loss += float(loss.detach())
batches += 1
if (epoch + 1) % 10 == 0 or epoch == 0:
print(f"{target_kind} epoch {epoch + 1}/{epochs}: loss={epoch_loss / batches:.4f}")
return model.eval()
def train_solution(
train_path: str | Path,
output_dir: str | Path = ".",
model_name: str = HERBERT_MODEL,
epochs: int = 40,
) -> tuple[DistanceModel, DepthModel, object, nn.Module]:
"""Train the task heads using only the organizer's labeled train set."""
torch.manual_seed(42)
np.random.seed(42)
torch.set_num_threads(min(4, torch.get_num_threads()))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(model_name)
bert = AutoModel.from_pretrained(model_name).to(device).eval()
sentences = read_conll(train_path)
embeddings = get_word_embeddings(sentences, tokenizer, bert)
distance_targets = [get_distances(sentence) for sentence in sentences]
depth_targets = [distance[sentence.root, :, None] for sentence, distance in zip(sentences, distance_targets)]
distance_model = DistanceModel(bert.config.hidden_size).to(device)
depth_model = DepthModel(bert.config.hidden_size).to(device)
_fit_regressor(distance_model, embeddings, distance_targets, "distance", epochs, 1e-3)
_fit_regressor(depth_model, embeddings, depth_targets, "depth", epochs, 1e-3)
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
torch.save(distance_model.cpu().state_dict(), output_dir / DISTANCE_MODEL_PATH)
torch.save(depth_model.cpu().state_dict(), output_dir / DEPTH_MODEL_PATH)
return distance_model, depth_model, tokenizer, bert.cpu()
def _minimum_spanning_tree(distances: np.ndarray) -> list[tuple[int, int]]:
n = distances.shape[0]
if n <= 1:
return []
weights = 0.5 * (distances + distances.T)
np.fill_diagonal(weights, np.inf)
selected = {0}
edges: list[tuple[int, int]] = []
while len(selected) < n:
best = (float("inf"), -1, -1)
for i in selected:
candidates = np.flatnonzero(~np.isin(np.arange(n), list(selected)))
if len(candidates) == 0:
break
j = int(candidates[np.argmin(weights[i, candidates])])
candidate = (float(weights[i, j]), i, j)
if candidate[0] < best[0]:
best = candidate
if best[1] < 0 or not np.isfinite(best[0]):
raise ValueError("Could not construct a spanning tree")
_, i, j = best
edges.append((i, j))
selected.add(j)
return edges
def parse_sentence(
sentence: Sentence,
distance_model: DistanceModel,
depth_model: DepthModel,
tokenizer,
bert: nn.Module,
) -> ParsedSentence:
embeddings = get_word_embeddings([sentence], tokenizer, bert)[0]
device = next(distance_model.parameters()).device
x = embeddings.unsqueeze(0).to(device)
with torch.no_grad():
pair_distances = distance_model(x)[0].cpu().numpy()
depths = depth_model(x)[0, :, 0].cpu().numpy()
root = int(np.argmin(depths))
edges = _minimum_spanning_tree(pair_distances)
return ParsedSentence.from_edges_and_root(sentence.words, edges, root)
def points(root_placement: float, uuas: float) -> float:
def scale(value: float) -> float:
return (min(max(value, 0.5), 0.85) - 0.5) / 0.35
return scale(root_placement) + scale(uuas)
def evaluate(
sentences: Sequence[ParsedSentence],
distance_model: DistanceModel,
depth_model: DepthModel,
tokenizer,
bert: nn.Module,
) -> tuple[float, float, float]:
root_correct = 0
uuas_total = 0.0
for sentence in sentences:
predicted = parse_sentence(sentence, distance_model, depth_model, tokenizer, bert)
if len(predicted.get_sorted_edges()) != len(sentence) - 1:
raise ValueError("Parser did not return a tree")
root_correct += int(predicted.root == sentence.root)
uuas_total += uuas_score(sentence, predicted)
root_accuracy = root_correct / len(sentences)
uuas = uuas_total / len(sentences)
score = points(root_accuracy, uuas)
print(f"Root placement: {root_accuracy:.4f}; UUAS: {uuas:.4f}; score: {score:.4f}/2")
return root_accuracy, uuas, score
def _load_head(model_class, path: Path, hidden_size: int, device):
model = model_class(hidden_size).to(device)
try:
state = torch.load(path, map_location=device, weights_only=True)
except TypeError:
state = torch.load(path, map_location=device)
model.load_state_dict(state)
return model.eval()
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--data-dir", type=Path, default=Path("."))
parser.add_argument("--model-name", default=HERBERT_MODEL)
parser.add_argument("--output-dir", type=Path, default=Path("."))
parser.add_argument("--epochs", type=int, default=40)
parser.add_argument("--train", action="store_true", help="Fit heads on train.conll first")
args = parser.parse_args()
torch.set_num_threads(min(4, torch.get_num_threads()))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
bert = AutoModel.from_pretrained(args.model_name).to(device).eval()
if args.train:
train_solution(
args.data_dir / "train.conll", args.output_dir, args.model_name, args.epochs
)
distance_model = _load_head(
DistanceModel, args.output_dir / DISTANCE_MODEL_PATH, bert.config.hidden_size, device
)
depth_model = _load_head(
DepthModel, args.output_dir / DEPTH_MODEL_PATH, bert.config.hidden_size, device
)
evaluate(read_conll(args.data_dir / "valid.conll"), distance_model, depth_model, tokenizer, bert)
if __name__ == "__main__":
main()Editorials are written by members of the SOTA community, not by the olympiad. The official statement and scoring win wherever an editorial differs.