albert-base-v2-persian-dep-parser

ALBERT Persian Dependency Parser

Syntactic dependency parser for Persian. A deep biaffine attention head (Dozat & Manning, 2017) on top of the shekar-ai/albert-base-v2-persian-zwnj-naab-mlm encoder, trained on the Persian Universal Dependency Treebank (PerDT). For each word it predicts a head (1-indexed; 0 = ROOT) and a Universal Dependencies relation label.

Files

The biaffine head is not a transformers architecture, so AutoModel cannot load the checkpoint on its own — the module below defines it, and the encoder inside it is a standard ALBERT.

Usage

Shekar

from shekar import DependencyParser

parser = DependencyParser()
result = parser("ما با آنچه می‌سازیم ایرانی هستیم.")

for word, head, relation in result:
    print(f"{word} ← (head: {head}, relation: {relation})")

ما ← (head: 6, relation: nsubj)
با ← (head: 3, relation: case)
آنچه ← (head: 6, relation: obl)
می‌سازیم ← (head: 3, relation: acl)
ایرانی ← (head: 6, relation: xcomp)
هستیم ← (head: 0, relation: root)
. ← (head: 6, relation: punct)

print_tree() renders the same parse as a tree:

parser.print_tree(result)

ROOT
└── [root] هستیم
    ├── [nsubj] ما
    ├── [obl] آنچه
    │   ├── [case] با
    │   └── [acl] می‌سازیم
    ├── [xcomp] ایرانی
    └── [punct] .

PyTorch

albert_persian_dep_parser.pt is a checkpoint dict holding model_state_dict alongside the dep_relations / deprel2id / id2deprel label maps, so the relation inventory comes from the file rather than being hard-coded. The forward pass takes input_ids, attention_mask, and word_to_first_subtoken (the index of each word's first subtoken), and returns arc_logits of shape (1, n+1, n+1) and rel_logits of shape (1, n+1, n+1, 34) — index 0 along the word axis is the ROOT slot.

import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel, AutoTokenizer

ENCODER = "shekar-ai/albert-base-v2-persian-zwnj-naab-mlm"


class BiaffineAttention(nn.Module):
    def __init__(self, dep_dim, head_dim, out_features=1):
        super().__init__()
        self.W = nn.Parameter(torch.zeros(out_features, dep_dim, head_dim))
        self.b_dep = nn.Parameter(torch.zeros(out_features, dep_dim))
        self.b_head = nn.Parameter(torch.zeros(out_features, head_dim))
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, dep, head):
        scores = torch.einsum("bnd,odh,bmh->bonm", dep, self.W, head)
        scores = scores + torch.einsum("bnd,od->bon", dep, self.b_dep).unsqueeze(-1)
        scores = scores + torch.einsum("bmh,oh->bom", head, self.b_head).unsqueeze(-2)
        return scores + self.bias.view(1, -1, 1, 1)


class BiaffineDependencyParser(nn.Module):
    def __init__(self, encoder_name, num_deprels, arc_hidden=500, rel_hidden=100, dropout=0.33):
        super().__init__()
        self.encoder = AutoModel.from_config(AutoConfig.from_pretrained(encoder_name))
        hidden_size = self.encoder.config.hidden_size

        def mlp(out_dim):
            return nn.Sequential(
                nn.Linear(hidden_size, out_dim), nn.LeakyReLU(0.1), nn.Dropout(dropout)
            )

        self.root_embedding = nn.Parameter(torch.zeros(1, 1, hidden_size))
        self.mlp_arc_dep, self.mlp_arc_head = mlp(arc_hidden), mlp(arc_hidden)
        self.mlp_rel_dep, self.mlp_rel_head = mlp(rel_hidden), mlp(rel_hidden)
        self.arc_biaffine = BiaffineAttention(arc_hidden, arc_hidden, 1)
        self.rel_biaffine = BiaffineAttention(rel_hidden, rel_hidden, num_deprels)

    def forward(self, input_ids, attention_mask, word_to_first_subtoken):
        hidden = self.encoder(
            input_ids=input_ids, attention_mask=attention_mask
        ).last_hidden_state

        index = word_to_first_subtoken.unsqueeze(-1).expand(-1, -1, hidden.size(-1))
        words = torch.gather(hidden, 1, index)
        words = torch.cat([self.root_embedding.expand(words.size(0), 1, -1), words], dim=1)

        arc_logits = self.arc_biaffine(self.mlp_arc_dep(words), self.mlp_arc_head(words)).squeeze(1)
        rel_logits = self.rel_biaffine(self.mlp_rel_dep(words), self.mlp_rel_head(words))
        return arc_logits, rel_logits.permute(0, 2, 3, 1)


checkpoint = torch.load(
    hf_hub_download("shekar-ai/albert-base-v2-persian-dep-parser", "albert_persian_dep_parser.pt"),
    map_location="cpu",
    weights_only=False,
)
id2deprel = checkpoint["id2deprel"]

tokenizer = AutoTokenizer.from_pretrained(ENCODER)
model = BiaffineDependencyParser(ENCODER, num_deprels=len(id2deprel))
# strict=False: the checkpoint carries a legacy `position_ids` buffer.
model.load_state_dict(checkpoint["model_state_dict"], strict=False)
model.eval()


@torch.no_grad()
def parse(words):
    input_ids, first_subtoken = [tokenizer.cls_token_id], []
    for word in words:
        subtokens = tokenizer.encode(word, add_special_tokens=False) or [tokenizer.unk_token_id]
        first_subtoken.append(len(input_ids))
        input_ids.extend(subtokens)
    input_ids.append(tokenizer.sep_token_id)

    arc_logits, rel_logits = model(
        input_ids=torch.tensor([input_ids]),
        attention_mask=torch.ones(1, len(input_ids), dtype=torch.long),
        word_to_first_subtoken=torch.tensor([first_subtoken]),
    )

    parsed = []
    for i in range(len(words)):
        head = int(arc_logits[0, i + 1, : len(words) + 1].argmax())
        parsed.append((words[i], head, id2deprel[int(rel_logits[0, i + 1, head].argmax())]))
    return parsed


for word, head, relation in parse(["ما", "با", "آنچه", "می‌سازیم", "ایرانی", "هستیم", "."]):
    print(f"{word} -> head={head} ({relation})")

ما -> head=6 (nsubj)
با -> head=3 (case)
آنچه -> head=6 (obl)
می‌سازیم -> head=3 (acl)
ایرانی -> head=6 (xcomp)
هستیم -> head=0 (root)
. -> head=6 (punct)

Training

Data. PerDT (commit d728a98): 26,196 training, 1,456 development, and 1,455 test sentences. Sentences longer than 150 words were dropped. Each word is tokenized independently and represented by its first subtoken; a learned ROOT embedding is prepended.

Architecture. Four LeakyReLU MLPs project word representations into arc-dependent, arc-head (500-dim), relation-dependent, and relation-head (100-dim) spaces. A biaffine layer scores every head candidate for each word; a second biaffine layer scores the 34 relations for each (dependent, head) pair. Training minimizes cross-entropy over gold heads plus cross-entropy over gold relations at the gold head.

Evaluation

On the PerDT test split (1,455 sentences, 24,133 tokens), using gold word segmentation and greedy head selection:

Excluding punctuation, the quantized model reaches 95.37 UAS / 92.42 LAS. Best development LAS during training was 91.69. Int8 quantization costs nothing measurable in accuracy here while cutting the graph to 37 MB.

Citation

@article{Amirivojdan2025Shekar,
  author  = {Amirivojdan, Ahmad},
  title   = {{Shekar: A Python Toolkit for Persian Natural Language Processing}},
  journal = {Journal of Open Source Software},
  volume  = {10},
  number  = {114},
  pages   = {9128},
  year    = {2025},
  doi     = {10.21105/joss.09128},
  url     = {https://joss.theoj.org/papers/10.21105/joss.09128}
}