XLM-RoBERTa-BiLSTM-CRF-Joint

XLM-RoBERTa-BiLSTM-CRF-Joint: Myanmar Word Segmentation + POS Tagging

A joint model for Myanmar (Burmese) Word Segmentation and Part-of-Speech Tagging using a custom architecture built on top of xlm-roberta-base with an Asymmetric BiLSTM and Dual CRF heads.

Model Architecture

XLM-RoBERTa (768) + Position Embedding (64)
    ↓
Asymmetric BiLSTM
  Forward LSTM  : hidden=256
  Backward LSTM : hidden=512
  Concatenated  : 768
    ├── Head 1 – Word Segmentation CRF  (4 BIES labels: B, I, E, S)
    └── Head 2 – POS Tagging CRF        (68 labels: BIES × 17 UPOS tags)

Key design choices:

• Position embedding encodes distance-from-end-of-sentence signals (dim=64)
• Asymmetric BiLSTM: forward hidden=256, backward hidden=512 to capture stronger right-context for Myanmar syllable boundary detection
• Dual CRF heads: separate CRF decoders for WS and POS to enforce label sequence constraints
• Joint loss: loss = CRF_WS + CRF_POS + 0.3 * (CE_WS + CE_POS)
• WeightedRandomSampler with 4x boost for rare POS tags (AUX, INTJ, SYM, PROPN, SCONJ, DET, X)
• Trained with AMP (FP16 Mixed Precision) and DataParallel (2x Tesla T4)

Dataset

• Split: 80/10/10 (random_state=42)
• Input granularity: Syllable-level (word-level CoNLL-U → syllable BIES conversion)
• POS tagset: Universal Dependencies UPOS (17 tags)

Labels

Word Segmentation (WS) — 4 labels

POS Tags — 17 UPOS tags (each with B/I/E/S prefix = 68 total labels)

ADJ, ADP, ADV, AUX, CCONJ, DET, INTJ, NOUN, NUM, PART, PRON, PROPN, PUNCT, SCONJ, SYM, VERB, X

Training Details

Evaluation Results

Test Set Performance

Per-class POS F1 (Test Set)

Repository Files

Usage

import torch
import json
from transformers import AutoTokenizer

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")

# Load label mappings
with open("ws_id2label.json") as f:
    ws_id2label = {int(k): v for k, v in json.load(f).items()}
with open("pos_id2label.json") as f:
    pos_id2label = {int(k): v for k, v in json.load(f).items()}

# Rebuild model (same architecture as training)
# ... (instantiate JointSegPosModel with same params)

# Load weights
model.load_state_dict(torch.load("best_model.pt", map_location="cpu"))
model.eval()

# Inference: input is a list of Myanmar syllables
syllables = ["သ", "ာ", "း", "အေ", "လ", "ော", "ကျ", "န်း"]
encoding = tokenizer(
    syllables,
    is_split_into_words=True,
    return_tensors="pt",
    truncation=True,
    max_length=128,
    padding="max_length"
)

with torch.no_grad():
    ws_preds, pos_preds = model(encoding["input_ids"], encoding["attention_mask"])

Citation

If you use this model, please cite:

@misc{sithu015_xlm_roberta_bilstm_crf_joint_2026,
  title  = {XLM-RoBERTa-BiLSTM-CRF-Joint: Myanmar Word Segmentation and POS Tagging},
  author = {sithu015},
  year   = {2026},
  url    = {https://huggingface.co/sithu015/XLM-RoBERTa-BiLSTM-CRF-Joint}
}

License

Apache 2.0

Acknowledgements

• Base model: xlm-roberta-base by Facebook AI
• Dataset: Myanmar CoNLL-U corpus (103,833 sentences)
• Training platform: Kaggle (2x Tesla T4 GPU)
• Libraries: transformers, pytorch-crf, seqeval, torch