Add PhayaThaiBERT POS tagger (ONNX and Transformers)

CHANGELOG.md

@@ -27,6 +27,8 @@ and this project adheres to
   ([#1499], [#1509])
 - `evaluate_word_tokenization`, `word_eval_function`, `char_eval_function`,
   and `TokenizationScore` in `pythainlp.benchmarks` ([#1512])
+- PhayaThaiBERT POS tagger: `pos_tag()` engine `phayathaibert` (ONNX) and
+  `pos_tag_transformers()` corpus `tud` ([#1550])
 - Laya Multilingual model (`LayaModel`, `Laya`) to `pythainlp.classify`
   for zero-shot classification ([#1528])
 - `CRFTagger`, CRF weight conversion, and bundled SSG weights ([#1538])
@@ -126,6 +128,7 @@ and this project adheres to
 [#1546]: https://github.com/PyThaiNLP/pythainlp/pull/1546
 [#1547]: https://github.com/PyThaiNLP/pythainlp/pull/1547
 [#1548]: https://github.com/PyThaiNLP/pythainlp/pull/1548
+[#1550]: https://github.com/PyThaiNLP/pythainlp/pull/1550
 [#1551]: https://github.com/PyThaiNLP/pythainlp/pull/1551
 [#1552]: https://github.com/PyThaiNLP/pythainlp/pull/1552
 [#1565]: https://github.com/PyThaiNLP/pythainlp/pull/1565

docs/api/tag.rst

@@ -228,6 +228,7 @@ Modules
 
 .. autofunction:: pos_tag
 .. autofunction:: pos_tag_sents
+.. autofunction:: pos_tag_transformers
 .. autofunction:: tag_provinces
 .. autofunction:: chunk_parse
 .. autoclass:: CRFTagger
@@ -253,6 +254,23 @@ unigram
 
 Unigram tagger doesn't take the ordering of words in the list into account.
 
+phayathaibert
++++++++++++++
+
+PhayaThaiBERT tagger [#Sriwirote_2025]_ uses
+`nlp-chula/phayathaibert-thai-pos-tagger <https://huggingface.co/nlp-chula/phayathaibert-thai-pos-tagger>`_,
+fine-tuned on the TUD treebank, and tags words with Universal POS tags.
+
+- :func:`pos_tag` with ``engine="phayathaibert"`` runs an ONNX export of the
+  model with ONNX Runtime and does not need PyTorch.
+  Install the dependencies with ``pip install "pythainlp[phayathaibert_onnx]"``.
+  The model (about 530 MB) is downloaded from the Hugging Face Hub on first use.
+- :func:`pos_tag_transformers` with ``engine="phayathaibert"`` and
+  ``corpus="tud"`` runs the original model with Transformers.
+
+.. autoclass:: pythainlp.tag.phayathaibert_onnx.PhayaThaiBERTTagger
+   :members: tag
+
 
 References
 ----------
@@ -263,3 +281,7 @@ References
 .. [#Prachya_2020] Prachya Boonkwan and Vorapon Luantangsrisuk and Sitthaa Phaholphinyo and Kanyanat Kriengket and Dhanon Leenoi and Charun Phrombut and Monthika Boriboon and Krit Kosawat and Thepchai Supnithi. (2020).
             The Annotation Guideline of LST20 Corpus.
             arXiv:2008.05055
+.. [#Sriwirote_2025] Panyut Sriwirote, Attapol T. Rutherford, Jalinee Thapiang and Vasan Timtong. (2025).
+            PhayaThaiBERT: Enhancing a Pretrained Thai Language Model with Unassimilated Loanwords.
+            ACM Transactions on Asian and Low-Resource Language Information Processing, Vol. 24, No. 11, pp 1-17.
+            doi:10.1145/3765962

docs/notes/installation.rst

@@ -32,6 +32,7 @@ The extras can include:
   - ``nlpo3`` — nlpo3 Thai word tokenization support
   - ``onnx`` - ONNX model support
   - ``oskut`` — OSKUT support (via LEKCut ONNX)
+  - ``phayathaibert_onnx`` — PhayaThaiBERT part-of-speech tagger with ONNX Runtime
   - ``sefr_cut`` — SEFR CUT Thai word tokenization support (via LEKCut ONNX)
   - ``spacy_thai`` — spaCy Thai language support
   - ``spell`` — support for more spell-checkers (phunspell & symspellpy)

pyproject.toml

@@ -149,6 +149,13 @@ onnx = ["numpy>=1.22", "onnxruntime>=1.10.0", "sentencepiece>=0.1.91"]
 
 oskut = ["lekcut>=1.0.0"]
 
+phayathaibert_onnx = [
+    "huggingface-hub>=0.16.0",
+    "numpy>=1.22",
+    "onnxruntime>=1.12.0",
+    "tokenizers>=0.19.0",
+]
+
 qwen3 = ["torch>=1.9.0", "transformers>=4.22.1"]
 
 sefr_cut = ["lekcut>=1.0.0"]

pythainlp/tag/phayathaibert_onnx.py

@@ -0,0 +1,266 @@
+# SPDX-FileCopyrightText: 2016-2026 PyThaiNLP Project
+# SPDX-FileType: SOURCE
+# SPDX-License-Identifier: Apache-2.0
+"""PhayaThaiBERT part-of-speech tagger with ONNX Runtime backend."""
+
+from __future__ import annotations
+
+import json
+import threading
+import unicodedata
+from collections import Counter
+from importlib.util import find_spec
+from typing import TYPE_CHECKING, Optional
+
+from pythainlp.corpus import get_hf_hub
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from onnxruntime import InferenceSession
+    from tokenizers import Encoding, Tokenizer
+
+_REPO_ID = "wiriyabot/phayathaibert-thai-pos-tagger-onnx"
+# Pinned for reproducible and secure downloads.
+_REVISION = "59e3502bd06e2222d32f11b51b4fb9569235a020"
+
+# Maximum number of subword tokens per model call, including <s> and </s>.
+# The model runs on up to 510 tokens, but its accuracy drops sharply from
+# position 415 on, which suggests it was pretrained on shorter sequences.
+_MAX_SEQUENCE_LENGTH = 415
+
+# Tag for words the model never sees (whitespace-only words), matching
+# the perceptron tagger with Universal POS corpora.
+_WHITESPACE_TAG = "PUNCT"
+
+# Universal POS tag for words that produce no subword token.
+_UNKNOWN_TAG = "X"
+
+# Modules the tagger needs at run time.
+_DEPENDENCIES = ("huggingface_hub", "numpy", "onnxruntime", "tokenizers")
+
+
+def _is_blank(word: str) -> bool:
+    """
+    Return True if a word has only whitespace or format characters.
+
+    Format characters (Unicode category Cf) include the zero-width space,
+    zero-width non-joiner and byte order mark, which carry no text.
+
+    :param str word: word to check
+    :return: whether the word is empty, whitespace or format characters
+    :rtype: bool
+    """
+    return all(ch.isspace() or unicodedata.category(ch) == "Cf" for ch in word)
+
+
+def _first_subword_labels(
+    word_ids: Sequence[Optional[int]],
+    label_ids: Sequence[int],
+    n_words: int,
+    id2label: Mapping[int, str],
+) -> list[str]:
+    """
+    Pick each word's tag from the prediction for its first subword.
+
+    This is the alignment the model was trained with.
+
+    :param Sequence[Optional[int]] word_ids: word index of each subword
+        token; ``None`` for special tokens
+    :param Sequence[int] label_ids: predicted label id of each subword token
+    :param int n_words: number of input words
+    :param Mapping[int, str] id2label: mapping from label id to tag name
+    :return: one tag per word; ``"X"`` for a word that produced no subword
+    :rtype: list[str]
+    """
+    tags = [_UNKNOWN_TAG] * n_words
+    seen: set[int] = set()
+    for word_id, label_id in zip(word_ids, label_ids):
+        if word_id is None or word_id in seen:
+            continue
+        seen.add(word_id)
+        tags[word_id] = id2label[label_id]
+    return tags
+
+
+def _chunk_spans(
+    subword_counts: Sequence[int], max_subwords: int
+) -> list[tuple[int, int]]:
+    """
+    Group consecutive words into spans that fit the model's input length.
+
+    A word longer than ``max_subwords`` gets a span of its own; its subwords
+    are truncated at encoding time, which keeps its first subword.
+
+    :param Sequence[int] subword_counts: number of subword tokens of each
+        word
+    :param int max_subwords: maximum number of subword tokens per span
+    :return: ``(start, end)`` word index pairs, end exclusive
+    :rtype: list[tuple[int, int]]
+    """
+    spans: list[tuple[int, int]] = []
+    start = 0
+    total = 0
+    for i, count in enumerate(subword_counts):
+        if i > start and total + count > max_subwords:
+            spans.append((start, i))
+            start = i
+            total = 0
+        total += count
+    if subword_counts:
+        spans.append((start, len(subword_counts)))
+    return spans
+
+
+class PhayaThaiBERTTagger:
+    """
+    Universal POS tagger using PhayaThaiBERT with ONNX Runtime.
+
+    The model is an ONNX export of
+    `nlp-chula/phayathaibert-thai-pos-tagger
+    <https://huggingface.co/nlp-chula/phayathaibert-thai-pos-tagger>`_,
+    fine-tuned on the UD Thai-TUD treebank.
+
+    Requires ``numpy``, ``onnxruntime``, ``tokenizers`` and
+    ``huggingface-hub``. The model (about 530 MB) is downloaded from
+    the Hugging Face Hub on first use.
+
+    :param str repo_id: Hugging Face Hub repository of the ONNX model
+    :param Optional[str] revision: git revision of the repository.
+        The default repository is pinned to a commit when this is ``None``.
+    """
+
+    session: InferenceSession
+    tokenizer: Tokenizer
+    id2label: dict[int, str]
+
+    def __init__(
+        self, repo_id: str = _REPO_ID, revision: Optional[str] = None
+    ) -> None:
+        """
+        Download the model and load it with ONNX Runtime.
+
+        :param str repo_id: Hugging Face Hub repository of the ONNX model
+        :param Optional[str] revision: git revision of the repository.
+            The default repository is pinned to a commit when this is
+            ``None``.
+        :raises ImportError: if a required dependency is not installed
+        """
+        # Check every dependency before the download starts.
+        if any(find_spec(name) is None for name in _DEPENDENCIES):
+            raise ImportError(
+                "PhayaThaiBERT POS tagger requires numpy, onnxruntime,"
+                " tokenizers and huggingface-hub."
+                ' Install them with: pip install "pythainlp[phayathaibert_onnx]"'
+            )
+        from onnxruntime import InferenceSession
+        from tokenizers import Tokenizer
+
+        if revision is None and repo_id == _REPO_ID:
+            revision = _REVISION
+        model_path = get_hf_hub(repo_id, "model.onnx", revision=revision)
+        tokenizer_path = get_hf_hub(
+            repo_id, "tokenizer.json", revision=revision
+        )
+        config_path = get_hf_hub(repo_id, "config.json", revision=revision)
+
+        self.session = InferenceSession(
+            model_path, providers=["CPUExecutionProvider"]
+        )
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+        # tokenizer.json ships with truncation at 510 tokens, which would
+        # silently drop words; chunking is done here instead.
+        self.tokenizer.no_truncation()
+        self.tokenizer.no_padding()
+        with open(config_path, encoding="utf-8") as f:
+            self.id2label = {
+                int(k): v for k, v in json.load(f)["id2label"].items()
+            }
+
+    def tag(self, words: list[str]) -> list[tuple[str, str]]:
+        """
+        Tag a list of words with Universal POS tags.
+
+        Words with only whitespace or format characters (such as the
+        zero-width space) are tagged ``PUNCT`` without being passed to the
+        model. Input longer than the model limit is tagged in consecutive
+        chunks of whole words.
+
+        :param list[str] words: list of words to be tagged
+        :return: list of tuples (word, POS tag)
+        :rtype: list[tuple[str, str]]
+        """
+        tags = [_WHITESPACE_TAG] * len(words)
+        positions = [i for i, word in enumerate(words) if not _is_blank(word)]
+        content = [words[i] for i in positions]
+        if content:
+            for position, word_tag in zip(
+                positions, self._tag_content(content)
+            ):
+                tags[position] = word_tag
+        return list(zip(words, tags))
+
+    def _tag_content(self, words: list[str]) -> list[str]:
+        encoding = self.tokenizer.encode(words, is_pretokenized=True)
+        if len(encoding.ids) <= _MAX_SEQUENCE_LENGTH:
+            return self._predict(encoding, len(words))
+
+        counts = Counter(i for i in encoding.word_ids if i is not None)
+        subword_counts = [counts.get(i, 0) for i in range(len(words))]
+        labels: list[str] = []
+        # Two positions are taken by <s> and </s>.
+        for start, end in _chunk_spans(
+            subword_counts, _MAX_SEQUENCE_LENGTH - 2
+        ):
+            chunk = words[start:end]
+            chunk_encoding = self.tokenizer.encode(chunk, is_pretokenized=True)
+            labels.extend(self._predict(chunk_encoding, len(chunk)))
+        return labels
+
+    def _predict(self, encoding: Encoding, n_words: int) -> list[str]:
+        import numpy as np
+
+        ids = encoding.ids
+        mask = encoding.attention_mask
+        word_ids = encoding.word_ids
+        if len(ids) > _MAX_SEQUENCE_LENGTH:
+            # Only a single over-long word gets here; its first subword stays.
+            ids = ids[: _MAX_SEQUENCE_LENGTH - 1] + ids[-1:]
+            mask = mask[:_MAX_SEQUENCE_LENGTH]
+            word_ids = [*word_ids[: _MAX_SEQUENCE_LENGTH - 1], None]
+        logits = self.session.run(
+            None,
+            {
+                "input_ids": np.array([ids], dtype=np.int64),
+                "attention_mask": np.array([mask], dtype=np.int64),
+            },
+        )[0][0]
+        return _first_subword_labels(
+            word_ids, logits.argmax(axis=-1).tolist(), n_words, self.id2label
+        )
+
+
+_TAGGER: Optional[PhayaThaiBERTTagger] = None
+_TAGGER_LOCK = threading.Lock()
+
+
+def tag(
+    words: list[str],
+    corpus: str = "tud",  # NOSONAR: unused, kept for the pos_tag interface
+) -> list[tuple[str, str]]:
+    """
+    Tag words with Universal POS tags using PhayaThaiBERT (ONNX).
+
+    :param list[str] words: list of words to be tagged
+    :param str corpus: ignored; the model is trained on UD Thai-TUD only
+    :return: list of tuples (word, POS tag)
+    :rtype: list[tuple[str, str]]
+    """
+    global _TAGGER
+    if not words:
+        return []
+    if _TAGGER is None:
+        with _TAGGER_LOCK:
+            if _TAGGER is None:
+                _TAGGER = PhayaThaiBERTTagger()
+    return _TAGGER.tag(words)