# SPDX-FileCopyrightText: 2016-2026 PyThaiNLP Project
# SPDX-FileType: SOURCE
# SPDX-License-Identifier: Apache-2.0
"""Provide text processing, augmentation, and tagging with PhayaThaiBERT."""
from __future__ import annotations
import random
import re
import warnings
from typing import TYPE_CHECKING, Optional, Union, cast
if TYPE_CHECKING:
from collections.abc import Callable
from transformers import (
AutoModelForMaskedLM,
AutoModelForTokenClassification,
CamembertTokenizer,
Pipeline,
PreTrainedTokenizerBase,
)
from pythainlp.tokenize import word_tokenize
_PAT_URL: str = (
r"(?:https?|ftp)://[\w-]+(?:\.[\w-]+)+"
r"(?:[\w.,@?^=%&:/~+#-]*[\w@?^=%&/~+#-])?"
)
_model_name: str = "clicknext/phayathaibert"
_tokenizer: Optional[CamembertTokenizer] = None
def _get_tokenizer() -> CamembertTokenizer:
"""Get the tokenizer, initializing it if necessary."""
global _tokenizer
if _tokenizer is None:
from transformers import CamembertTokenizer
_tokenizer = CamembertTokenizer.from_pretrained(
_model_name # nosec B615
)
return _tokenizer
[docs]
class ThaiTextProcessor:
"""Preprocess Thai text for PhayaThaiBERT."""
[docs]
def __init__(self) -> None:
"""Initialize the text processor."""
(
self._TK_UNK,
self._TK_REP,
self._TK_WREP,
self._TK_URL,
self._TK_END,
) = ["<unk>", "<rep>", "<wrep>", "<url>", "</s>"]
self.SPACE_SPECIAL_TOKEN: str = "<_>" # noqa: S105
[docs]
def replace_url(self, text: str) -> str:
"""
Replace URLs in text with the URL token.
See https://stackoverflow.com/a/6041965
:param str text: text to replace url
:return: text where urls are replaced
:rtype: str
:Example:
>>> replace_url("go to https://github.com")
'go to <url>'
"""
return re.sub(_PAT_URL, self._TK_URL, text)
[docs]
def rm_brackets(self, text: str) -> str:
"""
Remove empty brackets and artifacts within brackets from text.
:param str text: text to remove useless brackets
:return: text where all useless brackets are removed
:rtype: str
:Example:
>>> rm_brackets("hey() whats[;] up{*&} man(hey)")
'hey whats up man(hey)'
"""
# remove empty brackets
new_line = text.replace("()", "").replace("{}", "").replace("[]", "")
# brackets with only punctuations
new_line = re.sub(r"\([^a-zA-Z0-9ก-๙]+\)", "", new_line)
new_line = re.sub(r"\{[^a-zA-Z0-9ก-๙]+\}", "", new_line)
new_line = re.sub(r"\[[^a-zA-Z0-9ก-๙]+\]", "", new_line)
# artifiacts after (
new_line = re.sub(
r"(?<=\()[^a-zA-Z0-9ก-๙]+(?=[a-zA-Z0-9ก-๙])", "", new_line
)
new_line = re.sub(
r"(?<=\{)[^a-zA-Z0-9ก-๙]+(?=[a-zA-Z0-9ก-๙])", "", new_line
)
new_line = re.sub(
r"(?<=\[)[^a-zA-Z0-9ก-๙]+(?=[a-zA-Z0-9ก-๙])", "", new_line
)
# artifacts before )
new_line = re.sub(
r"(?<=[a-zA-Z0-9ก-๙])[^a-zA-Z0-9ก-๙]+(?=\))", "", new_line
)
new_line = re.sub(
r"(?<=[a-zA-Z0-9ก-๙])[^a-zA-Z0-9ก-๙]+(?=\})", "", new_line
)
new_line = re.sub(
r"(?<=[a-zA-Z0-9ก-๙])[^a-zA-Z0-9ก-๙]+(?=\])", "", new_line
)
return new_line
[docs]
def replace_newlines(self, text: str) -> str:
"""
Replace newlines in text with spaces.
:param str text: text to replace all newlines with spaces
:return: text where all newlines are replaced with spaces
:rtype: str
:Example:
>>> rm_useless_spaces("hey whats\n\nup")
hey whats up
"""
return text.strip().replace("\n", " ")
[docs]
def rm_useless_spaces(self, text: str) -> str:
"""
Collapse repeated spaces in text (code from `fastai`).
:param str text: text to replace useless spaces
:return: text where all spaces are reduced to one
:rtype: str
:Example:
>>> rm_useless_spaces("oh no")
oh no
"""
return re.sub(" {2,}", " ", text)
[docs]
def replace_spaces(self, text: str, space_token: str = "<_>") -> str: # noqa: S107 # nosec B107
"""
Replace spaces in text with a space token.
:param str text: text to replace spaces
:return: text where all spaces replaced with _
:rtype: str
:Example:
>>> replace_spaces("oh no")
oh_no
"""
return text.replace(" ", space_token)
[docs]
def replace_rep_after(self, text: str) -> str:
"""
Remove character repetitions in text.
:param str text: input text to replace character repetition
:return: text with repetitive tokens removed.
:rtype: str
:Example:
>>> text = "กาาาาาาา"
>>> replace_rep_after(text)
'กา'
"""
def _replace_rep(m: re.Match[str]) -> str:
c, _ = m.groups()
return f"{c}"
re_rep = re.compile(r"(\S)(\1{3,})")
return re_rep.sub(_replace_rep, text)
[docs]
def replace_wrep_post(self, toks: list[str]) -> list[str]:
"""
Remove repeated words after tokenization.
The `replace_wrep` function of `fastai` does not work well with Thai.
:param list[str] toks: list of tokens
:return: list of tokens where repetitive words are removed.
:rtype: list[str]
:Example:
>>> toks = ["กา", "น้ำ", "น้ำ", "น้ำ", "น้ำ"]
>>> replace_wrep_post(toks)
['กา', 'น้ำ']
"""
previous_word = ""
rep_count = 0
res = []
for current_word in [*toks, self._TK_END]:
if current_word == previous_word:
rep_count += 1
elif (current_word != previous_word) & (rep_count > 0):
res += [previous_word]
rep_count = 0
else:
res.append(previous_word)
previous_word = current_word
return res[1:]
[docs]
def remove_space(self, toks: list[str]) -> list[str]:
"""
Remove spaces from a list of words for bag-of-words models.
:param list[str] toks: list of tokens
:return: List of tokens where space tokens (" ") are filtered out
:rtype: list[str]
:Example:
>>> toks = ["ฉัน", "เดิน", " ", "กลับ", "บ้าน"]
>>> remove_space(toks)
['ฉัน', 'เดิน', 'กลับ', 'บ้าน']
"""
res = []
for token in toks:
stripped = token.strip()
if stripped:
res.append(stripped)
return res
# combine them together
[docs]
def preprocess(
self,
text: str,
pre_rules: Optional[list[Callable[..., str]]] = None,
tok_func: Callable[..., list[str]] = word_tokenize,
) -> str:
"""Apply preprocessing rules, tokenize text, and join the words."""
if pre_rules is None:
pre_rules = [
self.rm_brackets,
self.replace_newlines,
self.rm_useless_spaces,
self.replace_spaces,
self.replace_rep_after,
]
text = text.lower()
for rule in pre_rules:
text = rule(text)
toks = tok_func(text)
return "".join(toks)
[docs]
class ThaiTextAugmenter:
"""Generate and augment text with PhayaThaiBERT."""
[docs]
def __init__(self) -> None:
"""Initialize the text augmenter."""
from transformers import (
AutoModelForMaskedLM,
AutoTokenizer,
pipeline,
)
self.tokenizer: PreTrainedTokenizerBase = (
AutoTokenizer.from_pretrained(_model_name) # nosec B615
)
self.model_for_masked_lm: AutoModelForMaskedLM = (
AutoModelForMaskedLM.from_pretrained(_model_name) # nosec B615
)
self.model: Pipeline = pipeline( # transformers.Pipeline
"fill-mask",
tokenizer=self.tokenizer,
model=self.model_for_masked_lm,
)
self.processor: ThaiTextProcessor = ThaiTextProcessor()
[docs]
def generate(
self,
sample_text: str,
word_rank: int,
max_length: int = 3,
sample: bool = False,
) -> str:
"""Generate text from PhayaThaiBERT."""
sample_txt = sample_text
final_text = ""
for _ in range(max_length):
input_text = self.processor.preprocess(sample_txt)
if sample:
# Non-cryptographic use, pseudo-random generator is acceptable here
random_word_idx = random.randint(0, 4) # noqa: S311 # nosec B311 # NOSONAR
output = self.model(input_text)[random_word_idx]["sequence"]
else:
output = self.model(input_text)[word_rank]["sequence"]
sample_txt = output + "<mask>"
final_text = sample_txt
gen_txt = final_text.replace("<mask>", "")
return gen_txt
[docs]
def augment(
self,
text: str,
num_augs: int = 3,
sample: bool = False,
) -> list[str]:
"""
Text augmentation from PhayaThaiBERT.
:param str text: Thai text
:param int num_augs: an amount of augmentation text needed as an output
:param bool sample: whether to sample the text as an output or not,\
true if more word diversity is needed
:return: list of text augment
:rtype: list[str]
:Example:
>>> from pythainlp.augment.lm import ThaiTextAugmenter # doctest: +SKIP
>>> aug = ThaiTextAugmenter() # doctest: +SKIP
>>> aug.augment("ช้างมีทั้งหมด 50 ตัว บน", num_args=5) # doctest: +SKIP
['ช้างมีทั้งหมด 50 ตัว บนโลกใบนี้ครับ.',
'ช้างมีทั้งหมด 50 ตัว บนพื้นดินครับ...',
'ช้างมีทั้งหมด 50 ตัว บนท้องฟ้าครับ...',
'ช้างมีทั้งหมด 50 ตัว บนดวงจันทร์.‼',
'ช้างมีทั้งหมด 50 ตัว บนเขาค่ะ😁']
"""
MAX_NUM_AUGS = 5
augment_list = []
if num_augs <= MAX_NUM_AUGS:
for rank in range(num_augs):
gen_text = self.generate(
text,
rank,
sample=sample,
)
processed_text = self.processor.preprocess(gen_text).replace(
"<_>", " "
)
augment_list.append(processed_text)
else:
raise ValueError(
f"augmentation of more than {num_augs} is exceeded \
the default limit: {MAX_NUM_AUGS}"
)
return augment_list
[docs]
class PartOfSpeechTagger:
"""Tag Thai text with part-of-speech labels using PhayaThaiBERT."""
[docs]
def __init__(
self,
model: str = "lunarlist/pos_thai_phayathai",
revision: Optional[str] = None,
) -> None:
"""Initialize the part-of-speech tagger."""
# Load model directly
from transformers import (
AutoModelForTokenClassification,
AutoTokenizer,
)
self.tokenizer: PreTrainedTokenizerBase = (
AutoTokenizer.from_pretrained(model, revision=revision)
)
self.model: AutoModelForTokenClassification = (
AutoModelForTokenClassification.from_pretrained(
model, revision=revision
)
)
[docs]
def get_tag(
self, sentence: str, strategy: str = "simple"
) -> list[list[tuple[str, str]]]:
"""
Tag sentences with part-of-speech (POS) labels.
:param str sentence: a list of lists of tokenized words
:return: a list of lists of tuples (word, POS tag)
:rtype: list[list[tuple[str, str]]]
:Example:
Labels POS for given sentence:
>>> from pythainlp.lm.phayathaibert.core import (
... PartOfSpeechTagger,
... ) # doctest: +SKIP
>>> tagger = PartOfSpeechTagger() # doctest: +SKIP
>>> tagger.get_tag("แมวทำอะไรตอนห้าโมงเช้า") # doctest: +SKIP
[[('แมว', 'NOUN'), ('ทําอะไร', 'VERB'), ('ตอนห้าโมงเช้า', 'NOUN')]]
"""
from transformers import TokenClassificationPipeline
pipeline = TokenClassificationPipeline(
model=self.model,
tokenizer=self.tokenizer,
aggregation_strategy=strategy,
)
outputs = pipeline(sentence)
word_tags = [[(tag["word"], tag["entity_group"]) for tag in outputs]]
return word_tags
[docs]
class NamedEntityTagger:
"""Tag named entities with PhayaThaiBERT."""
[docs]
def __init__(
self,
model: str = "Pavarissy/phayathaibert-thainer",
revision: Optional[str] = None,
) -> None:
"""Initialize the named entity tagger."""
from transformers import (
AutoModelForTokenClassification,
AutoTokenizer,
)
self.tokenizer: PreTrainedTokenizerBase = (
AutoTokenizer.from_pretrained(model, revision=revision)
)
self.model: AutoModelForTokenClassification = (
AutoModelForTokenClassification.from_pretrained(
model, revision=revision
)
)
[docs]
def get_ner(
self,
text: str,
tag: bool = False,
pos: bool = False,
strategy: str = "simple",
) -> Union[list[tuple[str, str]], list[tuple[str, str, str]], str]:
"""
Tag named entities in text.
:param str text: text in Thai to be tagged
:param bool pos: output with part-of-speech tags.\
(PhayaThaiBERT is supported in PartOfSpeechTagger)
:return: a list of tuples associated with tokenized words, NER tags,
POS tags (if the parameter `pos` is specified as `True`),
and output HTML-like tags (if the parameter `tag` is
specified as `True`).
Otherwise, return a list of tuples associated with tokenized
words and NER tags
:rtype: Union[list[tuple[str, str]], list[tuple[str, str, str]], str]
:Example:
>>> from pythainlp.lm.phayathaibert.core import NamedEntityTagger
>>>
>>> tagger = NamedEntityTagger()
>>> tagger.get_ner("ทดสอบนายปวริศ เรืองจุติโพธิ์พานจากประเทศไทย")
[('นายปวริศ เรืองจุติโพธิ์พานจากประเทศไทย', 'PERSON'),
('จาก', 'LOCATION'),
('ประเทศไทย', 'LOCATION')]
>>> ner.tag("ทดสอบนายปวริศ เรืองจุติโพธิ์พานจากประเทศไทย", tag=True)
'ทดสอบ<PERSON>นายปวริศ เรืองจุติโพธิ์พาน</PERSON>\
<LOCATION>จาก</LOCATION><LOCATION>ประเทศไทย</LOCATION>'
"""
from transformers import TokenClassificationPipeline
if pos:
warnings.warn(
"This model does not support POS tag output.",
UserWarning,
stacklevel=2,
)
sample_output = []
tag_text_list = []
current_pos = 0
pipeline = TokenClassificationPipeline(
model=self.model,
tokenizer=self.tokenizer,
aggregation_strategy=strategy,
)
outputs = pipeline(text)
for token in outputs:
ner_tag = token["entity_group"]
begin_pos, end_pos = token["start"], token["end"]
if current_pos == 0:
text_tag = (
text[:begin_pos]
+ f"<{ner_tag}>"
+ text[begin_pos:end_pos]
+ f"</{ner_tag}>"
)
else:
text_tag = (
text[current_pos:begin_pos]
+ f"<{ner_tag}>"
+ text[begin_pos:end_pos]
+ f"</{ner_tag}>"
)
tag_text_list.append(text_tag)
sample_output.append((token["word"], token["entity_group"]))
current_pos = end_pos
if tag:
return str("".join(tag_text_list))
return sample_output
[docs]
def segment(sentence: str) -> list[str]:
"""
Tokenize text into subwords with the PhayaThaiBERT tokenizer.
The tokenizer is a SentencePiece model with vocabulary expansion.
:param str sentence: text to be tokenized
:return: list of subwords
:rtype: list[str]
"""
if not sentence or not isinstance(sentence, str):
return []
return cast("list[str]", _get_tokenizer().tokenize(sentence))