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 diff --git a/pythainlp/tag/named_entity.py b/pythainlp/tag/named_entity.py index 7a920f4c2..e768b789c 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,14 @@ ThaiNameTagger as WangchanbertaThaiNameTagger, ) + +class EntitySpan(TypedDict): + """Entity span dictionary""" + entity_type: str + text: list[str] + span: list[int] + + # Type alias for NER engine types NEREngineType = Union[ "ThaiNNER", @@ -164,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[dict[str, Any]]]: + ) -> tuple[list[str], list[EntitySpan]]: """This function tags nested named entities. :param str text: text in Thai to be tagged @@ -175,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[dict[str, Any]]] + :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 eb07506e2..1037ce226 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 EntitySpan 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: EntitySpan, container: EntitySpan) -> 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 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( def get_top_level_entities( - entities: list[dict[str, Any]], -) -> list[dict[str, Any]]: + 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[dict[str, Any]] 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[dict[str, Any]] + :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[dict[str, Any]] = [] + 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[dict]]: + ) -> 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[dict]] + :rtype: tuple[list[str], list[EntitySpan]] :Example: :: @@ -228,7 +228,7 @@ def get_ner( def _entities_to_iob( - tokens: list[str], entities: list[dict] + 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[dict] 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[dict]) -> 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[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[EntitySpan] entities: List of entity dictionaries :return: String with HTML-like entity tags :rtype: str """