From 66ccde40898f596f811ae32ff7e01d729a9ba7f0 Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 24 Mar 2026 00:10:43 +0000 Subject: [PATCH 1/3] Define NEREntity class for cleaner NER interface Instead of plain dict, use NEREntity with pre-defined keys ```python class NEREntity(TypedDict): text: list[str] span: list[int] entity_type: str ``` --- pythainlp/tag/named_entity.py | 13 ++++++++++--- pythainlp/tag/thai_nner.py | 34 +++++++++++++++++----------------- 2 files changed, 27 insertions(+), 20 deletions(-) diff --git a/pythainlp/tag/named_entity.py b/pythainlp/tag/named_entity.py index 7a920f4c2..d0b35226c 100644 --- a/pythainlp/tag/named_entity.py +++ b/pythainlp/tag/named_entity.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, TypedDict, Union if TYPE_CHECKING: from types import ModuleType @@ -20,6 +20,13 @@ ThaiNameTagger as WangchanbertaThaiNameTagger, ) + +class NEREntity(TypedDict): + entity_type: str + text: list[str] + span: list[int] + + # Type alias for NER engine types NEREngineType = Union[ "ThaiNNER", @@ -164,7 +171,7 @@ def load_engine(self, engine: str = "thai_nner") -> None: def tag( self, text: str, top_level_only: bool = False - ) -> tuple[list[str], list[dict[str, Any]]]: + ) -> tuple[list[str], list[NEREntity]]: """This function tags nested named entities. :param str text: text in Thai to be tagged @@ -175,7 +182,7 @@ def tag( :return: a tuple of (tokens, entities) where tokens is a list of tokenized strings and entities is a list of dictionaries containing 'text', 'span', and 'entity_type' keys. - :rtype: tuple[list[str], list[dict[str, Any]]] + :rtype: tuple[list[str], list[NEREntity]] .. note:: The tokenized output may include empty strings as part of the diff --git a/pythainlp/tag/thai_nner.py b/pythainlp/tag/thai_nner.py index eb07506e2..31748cdc2 100644 --- a/pythainlp/tag/thai_nner.py +++ b/pythainlp/tag/thai_nner.py @@ -9,23 +9,23 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Optional, Union from pythainlp.corpus import get_corpus_path +from pythainlp.tag.named_entity import NEREntity if TYPE_CHECKING: from thai_nner import NNER # noqa: F401 + __all__: list[str] = ["ThaiNNER"] -def _is_contained_in( - entity: dict[str, Any], container: dict[str, Any] -) -> bool: +def _is_contained_in(entity: NEREntity, container: NEREntity) -> bool: """Check if an entity is strictly contained within a container entity. - :param dict[str, Any] entity: Entity to check - :param dict[str, Any] container: Potential container entity + :param NEREntity entity: Entity to check + :param NEREntity container: Potential container entity :return: True if entity is strictly contained in container :rtype: bool """ @@ -42,19 +42,19 @@ def _is_contained_in( def get_top_level_entities( - entities: list[dict[str, Any]], -) -> list[dict[str, Any]]: + entities: list[NEREntity], +) -> list[NEREntity]: """Extract only top-level (outermost) entities from nested NER results. In nested NER, entities can contain other entities. This function filters the results to return only the outermost entities that are not contained within other entity. - :param list[dict[str, Any]] entities: List of entity dictionaries with + :param list[NEREntity] entities: List of entity dictionaries with 'span', 'text', and 'entity_type' keys :return: List of top-level entities only - :rtype: list[dict[str, Any]] + :rtype: list[NEREntity] :Example: :: @@ -81,7 +81,7 @@ def get_top_level_entities( entities, key=lambda x: (x["span"][0], -x["span"][1]) ) - top_level: list[dict[str, Any]] = [] + top_level: list[NEREntity] = [] for ent in sorted_entities: is_contained = False # Only check against entities already in top_level @@ -149,7 +149,7 @@ def __init__(self, path_model: Optional[str] = None) -> None: def tag( self, text: str, top_level_only: bool = False - ) -> tuple[list[str], list[dict]]: + ) -> tuple[list[str], list[NEREntity]]: """Tag Thai text with nested named entities. :param str text: Thai text to tag @@ -159,7 +159,7 @@ def tag( :return: Tuple of (tokens, entities) where tokens is a list of tokenized strings and entities is a list of dictionaries containing 'text', 'span', and 'entity_type' keys. - :rtype: tuple[list[str], list[dict]] + :rtype: tuple[list[str], list[NEREntity]] :Example: :: @@ -228,7 +228,7 @@ def get_ner( def _entities_to_iob( - tokens: list[str], entities: list[dict] + tokens: list[str], entities: list[NEREntity] ) -> list[tuple[str, str]]: """Convert Thai-NNER entity format to IOB format. @@ -238,7 +238,7 @@ def _entities_to_iob( will overwrite the IOB tags of earlier entities. :param list[str] tokens: List of tokens - :param list[dict] entities: List of entity dictionaries (should be non-overlapping) + :param list[NEREntity] entities: List of entity dictionaries (should be non-overlapping) :return: List of (token, tag) tuples in IOB format :rtype: list[tuple[str, str]] """ @@ -263,7 +263,7 @@ def _entities_to_iob( return result -def _entities_to_html(tokens: list[str], entities: list[dict]) -> str: +def _entities_to_html(tokens: list[str], entities: list[NEREntity]) -> str: """Convert Thai-NNER entity format to HTML-like tags. This function assumes entities do not overlap. If entities overlap, @@ -271,7 +271,7 @@ def _entities_to_html(tokens: list[str], entities: list[dict]) -> str: use only top-level entities (use get_top_level_entities() to filter). :param list[str] tokens: List of tokens - :param list[dict] entities: List of entity dictionaries + :param list[NEREntity] entities: List of entity dictionaries :return: String with HTML-like entity tags :rtype: str """ From ec7564558eb457488ed7a19423fb16c809220f93 Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 24 Mar 2026 00:17:36 +0000 Subject: [PATCH 2/3] Use EntitySpan instead of NEREntity --- pythainlp/tag/named_entity.py | 7 ++++--- pythainlp/tag/thai_nner.py | 30 +++++++++++++++--------------- 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/pythainlp/tag/named_entity.py b/pythainlp/tag/named_entity.py index d0b35226c..e768b789c 100644 --- a/pythainlp/tag/named_entity.py +++ b/pythainlp/tag/named_entity.py @@ -21,7 +21,8 @@ ) -class NEREntity(TypedDict): +class EntitySpan(TypedDict): + """Entity span dictionary""" entity_type: str text: list[str] span: list[int] @@ -171,7 +172,7 @@ def load_engine(self, engine: str = "thai_nner") -> None: def tag( self, text: str, top_level_only: bool = False - ) -> tuple[list[str], list[NEREntity]]: + ) -> tuple[list[str], list[EntitySpan]]: """This function tags nested named entities. :param str text: text in Thai to be tagged @@ -182,7 +183,7 @@ def tag( :return: a tuple of (tokens, entities) where tokens is a list of tokenized strings and entities is a list of dictionaries containing 'text', 'span', and 'entity_type' keys. - :rtype: tuple[list[str], list[NEREntity]] + :rtype: tuple[list[str], list[EntitySpan]] .. note:: The tokenized output may include empty strings as part of the diff --git a/pythainlp/tag/thai_nner.py b/pythainlp/tag/thai_nner.py index 31748cdc2..1037ce226 100644 --- a/pythainlp/tag/thai_nner.py +++ b/pythainlp/tag/thai_nner.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union from pythainlp.corpus import get_corpus_path -from pythainlp.tag.named_entity import NEREntity +from pythainlp.tag.named_entity import EntitySpan if TYPE_CHECKING: from thai_nner import NNER # noqa: F401 @@ -21,11 +21,11 @@ __all__: list[str] = ["ThaiNNER"] -def _is_contained_in(entity: NEREntity, container: NEREntity) -> bool: +def _is_contained_in(entity: EntitySpan, container: EntitySpan) -> bool: """Check if an entity is strictly contained within a container entity. - :param NEREntity entity: Entity to check - :param NEREntity container: Potential container entity + :param EntitySpan entity: Entity to check + :param EntitySpan container: Potential container entity :return: True if entity is strictly contained in container :rtype: bool """ @@ -42,19 +42,19 @@ def _is_contained_in(entity: NEREntity, container: NEREntity) -> bool: def get_top_level_entities( - entities: list[NEREntity], -) -> list[NEREntity]: + entities: list[EntitySpan], +) -> list[EntitySpan]: """Extract only top-level (outermost) entities from nested NER results. In nested NER, entities can contain other entities. This function filters the results to return only the outermost entities that are not contained within other entity. - :param list[NEREntity] entities: List of entity dictionaries with + :param list[EntitySpan] entities: List of entity dictionaries with 'span', 'text', and 'entity_type' keys :return: List of top-level entities only - :rtype: list[NEREntity] + :rtype: list[EntitySpan] :Example: :: @@ -81,7 +81,7 @@ def get_top_level_entities( entities, key=lambda x: (x["span"][0], -x["span"][1]) ) - top_level: list[NEREntity] = [] + top_level: list[EntitySpan] = [] for ent in sorted_entities: is_contained = False # Only check against entities already in top_level @@ -149,7 +149,7 @@ def __init__(self, path_model: Optional[str] = None) -> None: def tag( self, text: str, top_level_only: bool = False - ) -> tuple[list[str], list[NEREntity]]: + ) -> tuple[list[str], list[EntitySpan]]: """Tag Thai text with nested named entities. :param str text: Thai text to tag @@ -159,7 +159,7 @@ def tag( :return: Tuple of (tokens, entities) where tokens is a list of tokenized strings and entities is a list of dictionaries containing 'text', 'span', and 'entity_type' keys. - :rtype: tuple[list[str], list[NEREntity]] + :rtype: tuple[list[str], list[EntitySpan]] :Example: :: @@ -228,7 +228,7 @@ def get_ner( def _entities_to_iob( - tokens: list[str], entities: list[NEREntity] + tokens: list[str], entities: list[EntitySpan] ) -> list[tuple[str, str]]: """Convert Thai-NNER entity format to IOB format. @@ -238,7 +238,7 @@ def _entities_to_iob( will overwrite the IOB tags of earlier entities. :param list[str] tokens: List of tokens - :param list[NEREntity] entities: List of entity dictionaries (should be non-overlapping) + :param list[EntitySpan] entities: List of entity dictionaries (should be non-overlapping) :return: List of (token, tag) tuples in IOB format :rtype: list[tuple[str, str]] """ @@ -263,7 +263,7 @@ def _entities_to_iob( return result -def _entities_to_html(tokens: list[str], entities: list[NEREntity]) -> str: +def _entities_to_html(tokens: list[str], entities: list[EntitySpan]) -> str: """Convert Thai-NNER entity format to HTML-like tags. This function assumes entities do not overlap. If entities overlap, @@ -271,7 +271,7 @@ def _entities_to_html(tokens: list[str], entities: list[NEREntity]) -> str: use only top-level entities (use get_top_level_entities() to filter). :param list[str] tokens: List of tokens - :param list[NEREntity] entities: List of entity dictionaries + :param list[EntitySpan] entities: List of entity dictionaries :return: String with HTML-like entity tags :rtype: str """ From 83d8e67d835ea1794123b726125c1c0da8c9579a Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 24 Mar 2026 00:17:54 +0000 Subject: [PATCH 3/3] Update pyproject.toml --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index dbaac9f9a..ebb62207b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -414,7 +414,7 @@ basepython = "pypy3.11" [tool.tox.env.ruff] basepython = "python" -deps = ["ruff"] +deps = "ruff" commands = [["ruff", "check", "pythainlp"], ["ruff", "format", "--check", "pythainlp"]] skip_install = true