WangchanBERTa-NER-v0

WangchanBERTa fine-tuned on Thai NER v2.2

This is a Thai Named Entity Recognition (NER) model fine-tuned from airesearch/wangchanberta-base-att-spm-uncased on the Thai NER v2.2 corpus.

The base WangchanBERTa checkpoint was used (pretrained MLM only — not the existing fine-tuned variant) and a fresh token-classification head was trained from scratch.

Model details

• Base model: airesearch/wangchanberta-base-att-spm-uncased (RoBERTa-base, ~105M params, SentencePiece tokenizer, vocab 25k)
• Task: Token classification (NER) in IOB format
• Language: Thai
• Max sequence length: 416 subword tokens (matches pretraining)
• License: CC-BY-SA 3.0 (inherits from the training corpus)

Intended use

Entity extraction from general-domain Thai text (news, PR content, social media). The model predicts 14 entity types plus O, in IOB scheme (29 labels total).

Entity types

Training data

pythainlp/thainer-corpus-v2.2 — 7,326 documents (news, PR from KKU, general Thai text), pre-split by the dataset authors:

Tags in the corpus are stored as integers without a ClassLabel feature; the canonical id2label mapping was imported from pythainlp/thainer-corpus-v2-base-model which was trained on the same corpus.

Training procedure

Hyperparameters follow the WangchanBERTa paper's downstream recipe.

Preprocessing notes

Two WangchanBERTa-specific steps matter for correct results:

1. Space handling. Spaces are semantically meaningful in Thai. Any whitespace-only token in the pre-tokenized input is replaced with the model's space token <_> before passing to the tokenizer.
2. Label alignment. After subword tokenization, the word's tag is assigned to the first subword only; continuation subwords and special tokens are masked with -100 so they are ignored by the loss.

Hardware

Single premium GPU on Google Colab Pro+ RTX PRO 6000 Blackwell Server Edition.

Evaluation

Entity-level metrics computed with seqeval on the held-out test split (1,313 docs).

# Per-entity classification report on the test set
              precision    recall  f1-score   support

         AGO     0.7292    0.9211    0.8140        38
        DATE     0.8122    0.9276    0.8661       373
       EMAIL     1.0000    1.0000    1.0000         1
    FACILITY     0.4214    0.4558    0.4379       147
         LAW     0.5000    0.5909    0.5417        44
         LEN     1.0000    0.9412    0.9697        17
    LOCATION     0.8159    0.8558    0.8354       839
       MONEY     0.7800    0.8797    0.8269       133
ORGANIZATION     0.8453    0.9144    0.8785      1285
     PERCENT     0.6800    0.7907    0.7312        43
      PERSON     0.8873    0.9395    0.9127       645
       PHONE     0.9459    1.0000    0.9722        35
 TEMPERATURE     0.0000    0.0000    0.0000         2
        TIME     0.7436    0.7838    0.7632       185
         URL     0.9583    1.0000    0.9787        23
         ZIP     1.0000    1.0000    1.0000         6

   micro avg     0.8137    0.8779    0.8446      3816
   macro avg     0.7574    0.8125    0.7830      3816
weighted avg     0.8143    0.8779    0.8447      3816

How to use

Because of how WangchanBERTa handles spaces, the standard pipeline("ner", ...) gives incorrect results on Thai. Use PyThaiNLP for word-level tokenization and run the model manually:

import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from pythainlp.tokenize import word_tokenize

MODEL = "1amBanana/WangchanBERTa-NER-v0"
SPACE = "<_>"

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForTokenClassification.from_pretrained(MODEL).eval()
id2label = model.config.id2label

def predict_ner(text: str):
    words = word_tokenize(text, engine="newmm", keep_whitespace=True)
    norm = [SPACE if (w is None or not w.strip()) else w for w in words]
    enc = tokenizer(norm, is_split_into_words=True, return_tensors="pt",
                    truncation=True, max_length=416)
    with torch.no_grad():
        logits = model(**enc).logits[0]
    pred_ids = logits.argmax(-1).tolist()
    word_ids = enc.word_ids(batch_index=0)
    word_tags = {}
    for p, w in zip(pred_ids, word_ids):
        if w is None or w in word_tags:
            continue
        word_tags[w] = id2label[int(p)]
    return [(words[i], word_tags.get(i, "O")) for i in range(len(words))]

text = "นายมะลิวาโย บุญรักษา อาศัยอยู่ที่อำเภอนางรอง จังหวัดบุรีรัมย์ อายุ 99 ปี เพิ่งเรียนจบจากมหาวิทยาลัยเชียงใหม่"
result = predict_ner(sample)
pd.DataFrame(result, columns=["word", "tag"])

Here is the result of the block code:

Limitations

• Domain. The corpus is dominated by news and PR text. Performance on informal / social / domain-specific Thai (medical, legal, financial) will be lower.
• IOB scheme only. Non-nested entities. For overlapping entities, see Thai N-NER.
• <_> preprocessing. Required for correct inference. The default transformers NER pipeline does not do this automatically and will produce wrong tags.
• Case. The base model is uncased; casing is lost.
• Max length. Inputs longer than 416 subword tokens are truncated.

Citation

Base model:

@misc{lowphansirikul2021wangchanberta,
  title={WangchanBERTa: Pretraining transformer-based Thai Language Models},
  author={Lalita Lowphansirikul and Charin Polpanumas and Nawat Jantrakulchai and Sarana Nutanong},
  year={2021},
  eprint={2101.09635},
  archivePrefix={arXiv},
  primaryClass={cs.CL}
}

Dataset:

@dataset{wannaphong_phatthiyaphaibun_2024_10795907,
  author    = {Wannaphong Phatthiyaphaibun},
  title     = {Thai NER 2.2},
  month     = mar,
  year      = 2024,
  publisher = {Zenodo},
  version   = {2.2},
  doi       = {10.5281/zenodo.10795907},
  url       = {https://doi.org/10.5281/zenodo.10795907}
}