Merge pull request #1363 from bact/nerentity

Define EntitySpan class for cleaner NER interface

Author: bact

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
 

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

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,15 +263,15 @@ 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,
     tokens between overlapping entities may be skipped. For best results,
     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
     """