Add TypedDicts for complex return dicts: RougeScore, TokenizationStat family, CorefResult

CHANGELOG.md

@@ -19,6 +19,20 @@ and this project adheres to
 
 ## [Unreleased]
 
+### Added
+
+- `EntitySpan` TypedDict (#1363).
+  Migration notes:
+
+  ```python
+  # EntitySpan – Before (plain dict)
+  entity = {"text": ["สมชาย"], "span": [0, 1], "entity_type": "PERSON"}
+
+  # EntitySpan – After (TypedDict)
+  from pythainlp.tag.named_entity import EntitySpan
+  entity = EntitySpan(text=["สมชาย"], span=[0, 1], entity_type="PERSON")
+  ```
+
 ### Fixed
 
 - thai2rom_onnx: fix ONNX encoder model and fix inference bugs (#1349)

pythainlp/benchmarks/__init__.py

@@ -5,6 +5,11 @@
 
 __all__: list[str] = [
     "BleuScore",
+    "CharLevelStat",
+    "GlobalStat",
+    "RougeScore",
+    "TokenizationStat",
+    "WordLevelStat",
     "benchmark",
     "bleu_score",
     "character_error_rate",
@@ -14,9 +19,16 @@
 
 from pythainlp.benchmarks.metrics import (
     BleuScore,
+    RougeScore,
     bleu_score,
     character_error_rate,
     rouge_score,
     word_error_rate,
 )
-from pythainlp.benchmarks.word_tokenization import benchmark
+from pythainlp.benchmarks.word_tokenization import (
+    CharLevelStat,
+    GlobalStat,
+    TokenizationStat,
+    WordLevelStat,
+    benchmark,
+)

pythainlp/benchmarks/metrics.py

@@ -16,7 +16,7 @@
 
 
 class BleuScore(TypedDict):
-    """BLEU score"""
+    """BLEU score components returned by :func:`bleu_score`."""
 
     bleu: float  # BLEU score as a percentage (0.0 to 100.0)
     precisions: list[float]
@@ -26,6 +26,14 @@ class BleuScore(TypedDict):
     ref_length: int
 
 
+class RougeScore(TypedDict):
+    """Precision, recall, and F-measure for a single ROUGE type."""
+
+    precision: float
+    recall: float
+    fmeasure: float
+
+
 def _get_ngrams(tokens: list[str], n: int) -> list[tuple[str, ...]]:
     """
     Get n-grams from a list of tokens.
@@ -249,7 +257,7 @@ def rouge_score(
     hypothesis: str,
     tokenize: str = "newmm",
     rouge_types: Optional[list[str]] = None,
-) -> dict[str, tuple[float, float, float]]:
+) -> dict[str, RougeScore]:
     """
     Calculate ROUGE scores for Thai text with automatic tokenization.
 
@@ -269,8 +277,9 @@ def rouge_score(
     :param Optional[list[str]] rouge_types: list of ROUGE types to calculate.
         Default is ["rouge1", "rouge2", "rougeL"]
 
-    :return: dictionary mapping ROUGE type to (precision, recall, fmeasure)
-    :rtype: dict[str, tuple[float, float, float]]
+    :return: dictionary mapping ROUGE type to a :class:`RougeScore` typed dict
+        with ``'precision'``, ``'recall'``, and ``'fmeasure'`` keys.
+    :rtype: dict[str, RougeScore]
 
     :Example:
     ::
@@ -280,9 +289,9 @@ def rouge_score(
         reference = "สวัสดีครับ วันนี้อากาศดีมาก"
         hypothesis = "สวัสดีค่ะ วันนี้อากาศดี"
         scores = rouge_score(reference, hypothesis)
-        print(f"ROUGE-1 F-measure: {scores['rouge1'][2]:.4f}")
-        print(f"ROUGE-2 F-measure: {scores['rouge2'][2]:.4f}")
-        print(f"ROUGE-L F-measure: {scores['rougeL'][2]:.4f}")
+        print(f"ROUGE-1 F-measure: {scores['rouge1']['fmeasure']:.4f}")
+        print(f"ROUGE-2 F-measure: {scores['rouge2']['fmeasure']:.4f}")
+        print(f"ROUGE-L F-measure: {scores['rougeL']['fmeasure']:.4f}")
     """
     from pythainlp.tokenize import word_tokenize
 
@@ -297,7 +306,7 @@ def rouge_score(
         hypothesis, engine=tokenize, keep_whitespace=False
     )
 
-    result: dict[str, tuple[float, float, float]] = {}
+    result: dict[str, RougeScore] = {}
 
     for rouge_type in rouge_types:
         if rouge_type == "rouge1":
@@ -309,9 +318,12 @@ def rouge_score(
             ref_count = len(ref_tokens)
             hyp_count = len(hyp_tokens)
 
-            result[rouge_type] = _calculate_precision_recall_fmeasure(
+            precision, recall, fmeasure = _calculate_precision_recall_fmeasure(
                 overlap, hyp_count, ref_count
             )
+            result[rouge_type] = RougeScore(
+                precision=precision, recall=recall, fmeasure=fmeasure
+            )
 
         elif rouge_type == "rouge2":
             # Bigram-based
@@ -325,19 +337,25 @@ def rouge_score(
             ref_count = len(ref_bigrams)
             hyp_count = len(hyp_bigrams)
 
-            result[rouge_type] = _calculate_precision_recall_fmeasure(
+            precision, recall, fmeasure = _calculate_precision_recall_fmeasure(
                 overlap, hyp_count, ref_count
             )
+            result[rouge_type] = RougeScore(
+                precision=precision, recall=recall, fmeasure=fmeasure
+            )
 
         elif rouge_type == "rougeL":
             # Longest Common Subsequence-based
             lcs_len = _lcs_length(ref_tokens, hyp_tokens)
             ref_count = len(ref_tokens)
             hyp_count = len(hyp_tokens)
 
-            result[rouge_type] = _calculate_precision_recall_fmeasure(
+            precision, recall, fmeasure = _calculate_precision_recall_fmeasure(
                 lcs_len, hyp_count, ref_count
             )
+            result[rouge_type] = RougeScore(
+                precision=precision, recall=recall, fmeasure=fmeasure
+            )
 
     return result
 

pythainlp/benchmarks/word_tokenization.py

@@ -5,7 +5,8 @@
 
 import re
 import sys
-from typing import TYPE_CHECKING, Union
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, TypedDict, Union, overload
 
 if TYPE_CHECKING:
     import numpy as np
@@ -29,6 +30,37 @@
 TAILING_SEP_RX: re.Pattern[str] = re.compile(f"{re.escape(SEPARATOR)}$")
 
 
+class CharLevelStat(TypedDict):
+    """Character-level confusion matrix statistics for tokenization."""
+
+    tp: int
+    fp: int
+    tn: int
+    fn: int
+
+
+class WordLevelStat(TypedDict):
+    """Word-level tokenization statistics."""
+
+    correctly_tokenized_words: int
+    total_words_in_sample: int
+    total_words_in_ref_sample: int
+
+
+class GlobalStat(TypedDict):
+    """Global tokenization indicator as a binary indicator string."""
+
+    tokenization_indicators: str
+
+
+class TokenizationStat(TypedDict):
+    """Tokenization quality statistics at character, word, and global level."""
+
+    char_level: CharLevelStat
+    word_level: WordLevelStat
+    global_: GlobalStat
+
+
 def _f1(precision: float, recall: float) -> float:
     """Compute f1.
 
@@ -43,8 +75,21 @@ def _f1(precision: float, recall: float) -> float:
     return 2 * precision * recall / (precision + recall)
 
 
+@overload
+def _flatten_result(
+    my_dict: TokenizationStat, sep: str = ...
+) -> dict[str, Union[int, str]]: ...
+
+
+@overload
+def _flatten_result(
+    my_dict: Mapping[str, Mapping[str, Union[int, str]]], sep: str = ...
+) -> dict[str, Union[int, str]]: ...
+
+
 def _flatten_result(
-    my_dict: dict[str, dict[str, Union[int, str]]], sep: str = ":"
+    my_dict: Any,
+    sep: str = ":",
 ) -> dict[str, Union[int, str]]:
     """Flatten two-dimension dictionary.
 
@@ -55,8 +100,9 @@ def _flatten_result(
     { "a:b": 7 }
 
 
-    :param dict[str, dict[str, Union[int, str]]] my_dict: dictionary
-        containing stats
+    :param my_dict: dictionary containing stats
+    :type my_dict: TokenizationStat or
+        collections.abc.Mapping[str, collections.abc.Mapping[str, Union[int, str]]]
     :param str sep: separator between the two keys (default: ":")
 
     :return: a one-dimension dictionary with keys combined
@@ -139,7 +185,7 @@ def preprocessing(txt: str, remove_space: bool = True) -> str:
 
 def compute_stats(
     ref_sample: str, raw_sample: str
-) -> dict[str, dict[str, Union[int, str]]]:
+) -> TokenizationStat:
     """Compute statistics for tokenization quality
 
     These statistics include:
@@ -156,7 +202,7 @@ def compute_stats(
     :param str samples: samples that we want to evaluate
 
     :return: metrics at character- and word-level and indicators of correctly tokenized words
-    :rtype: dict[str, dict[str, Union[int, str]]]
+    :rtype: TokenizationStat
     """
     import numpy as np
 
@@ -185,29 +231,29 @@ def compute_stats(
 
     # Find correctly tokenized words in the sample
     ss_boundaries = _find_word_boundaries(sample_arr)
-    tokenization_indicators = _find_words_correctly_tokenised(
+    tokenization_indicators = _find_words_correctly_tokenized(
         word_boundaries, ss_boundaries
     )
 
-    correctly_tokenised_words: int = int(np.sum(tokenization_indicators))
+    correctly_tokenized_words: int = int(np.sum(tokenization_indicators))
 
     tokenization_indicators_str = list(map(str, tokenization_indicators))
 
     return {
-        "char_level": {
-            "tp": c_tp,
-            "fp": c_fp,
-            "tn": c_tn,
-            "fn": c_fn,
-        },
-        "word_level": {
-            "correctly_tokenised_words": correctly_tokenised_words,
-            "total_words_in_sample": int(np.sum(sample_arr)),
-            "total_words_in_ref_sample": int(np.sum(ref_sample_arr)),
-        },
-        "global": {
-            "tokenisation_indicators": "".join(tokenization_indicators_str)
-        },
+        "char_level": CharLevelStat(
+            tp=c_tp,
+            fp=c_fp,
+            tn=c_tn,
+            fn=c_fn,
+        ),
+        "word_level": WordLevelStat(
+            correctly_tokenized_words=correctly_tokenized_words,
+            total_words_in_sample=int(np.sum(sample_arr)),
+            total_words_in_ref_sample=int(np.sum(ref_sample_arr)),
+        ),
+        "global_": GlobalStat(
+            tokenization_indicators="".join(tokenization_indicators_str),
+        ),
     }
 
 
@@ -271,7 +317,7 @@ def _find_word_boundaries(
     return list(zip(start_idx, end_idx))
 
 
-def _find_words_correctly_tokenised(
+def _find_words_correctly_tokenized(
     ref_boundaries: list[tuple[int, int]],
     predicted_boundaries: list[tuple[int, int]],
 ) -> tuple[int, ...]:

pythainlp/cli/benchmark.py

@@ -108,7 +108,7 @@
             "char_level:fp",
             "char_level:tn",
             "char_level:fn",
-            "word_level:correctly_tokenised_words",
+            "word_level:correctly_tokenized_words",
             "word_level:total_words_in_sample",
             "word_level:total_words_in_ref_sample",
         ]
@@ -127,12 +127,12 @@
         )
 
         statistics["word_level:precision"] = (
-            statistics["word_level:correctly_tokenised_words"]
+            statistics["word_level:correctly_tokenized_words"]
             / statistics["word_level:total_words_in_sample"]
         )
 
         statistics["word_level:recall"] = (
-            statistics["word_level:correctly_tokenised_words"]
+            statistics["word_level:correctly_tokenized_words"]
             / statistics["word_level:total_words_in_ref_sample"]
         )
 
@@ -146,7 +146,7 @@
         for c in [
             "total_words_in_sample",
             "total_words_in_ref_sample",
-            "correctly_tokenised_words",
+            "correctly_tokenized_words",
             "precision",
             "recall",
         ]:

pythainlp/coref/__init__.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """PyThaiNLP Coreference Resolution"""
 
-__all__: list[str] = ["coreference_resolution"]
+__all__: list[str] = ["CorefResult", "coreference_resolution"]
 
+from pythainlp.coref._fastcoref import CorefResult
 from pythainlp.coref.core import coreference_resolution

pythainlp/coref/_fastcoref.py

@@ -6,12 +6,13 @@
 from typing import TYPE_CHECKING, Optional, TypedDict
 
 if TYPE_CHECKING:
-    from fastcoref.modeling import CorefModel, CorefResult
+    from fastcoref.modeling import CorefModel
+    from fastcoref.modeling import CorefResult as FastCorefResult
     from spacy.language import Language
 
 
-class CorefResultDict(TypedDict):
-    """Dictionary representation of coreference resolution results."""
+class CorefResult(TypedDict):
+    """Coreference resolution result for a single text."""
 
     text: str
     clusters_string: list[list[str]]
@@ -42,14 +43,14 @@ def __init__(
             self.model_name, device=device, nlp=self.nlp
         )
 
-    def _to_json(self, _predict: "CorefResult") -> CorefResultDict:
+    def _to_json(self, _predict: "FastCorefResult") -> CorefResult:
         return {
             "text": _predict.text,
             "clusters_string": _predict.get_clusters(as_strings=True),
             "clusters": _predict.get_clusters(as_strings=False),
         }
 
-    def predict(self, texts: list[str]) -> list[CorefResultDict]:
+    def predict(self, texts: list[str]) -> list[CorefResult]:
         return [
             self._to_json(pred) for pred in self.model.predict(texts=texts)
         ]