Use BleuScore TypeDict for BLEU score

pythainlp/benchmarks/__init__.py

@@ -4,6 +4,7 @@
 """Performance benchmarking."""
 
 __all__: list[str] = [
+    "BleuScore",
     "benchmark",
     "bleu_score",
     "character_error_rate",
@@ -12,6 +13,7 @@
 ]
 
 from pythainlp.benchmarks.metrics import (
+    BleuScore,
     bleu_score,
     character_error_rate,
     rouge_score,

pythainlp/benchmarks/metrics.py

@@ -12,7 +12,18 @@
 
 import math
 from collections import Counter
-from typing import Optional, Union, cast
+from typing import Optional, TypedDict, Union, cast
+
+
+class BleuScore(TypedDict):
+    """BLEU score"""
+
+    bleu: float  # BLEU score as a percentage (0.0 to 100.0)
+    precisions: list[float]
+    bp: float
+    length_ratio: float
+    hyp_length: int
+    ref_length: int
 
 
 def _get_ngrams(tokens: list[str], n: int) -> list[tuple[str, ...]]:
@@ -83,7 +94,7 @@ def bleu_score(
     lowercase: bool = False,
     max_ngram: int = 4,
     smooth: bool = True,
-) -> dict[str, Union[float, list[float]]]:
+) -> BleuScore:
     """
     Calculate BLEU score for Thai text with automatic tokenization.
 
@@ -104,11 +115,12 @@ def bleu_score(
     :param bool smooth: whether to use smoothing for zero counts
         (default: True)
 
-    :return: dictionary with ``'bleu'``, ``'precisions'``, ``'bp'``,
-        ``'length_ratio'``, ``'hyp_length'``, and ``'ref_length'``.
-        The ``'precisions'`` value is a ``list[float]``; all other values
-        are ``float``.
-    :rtype: dict[str, Union[float, list[float]]]
+    :return: a :class:`BleuScore` typed dict with ``'bleu'``,
+        ``'precisions'``, ``'bp'``, ``'length_ratio'``, ``'hyp_length'``,
+        and ``'ref_length'``. ``'precisions'`` is ``list[float]``;
+        ``'hyp_length'`` and ``'ref_length'`` are ``int``; all other
+        values are ``float``.
+    :rtype: BleuScore
 
     :Example:
     ::
@@ -220,16 +232,16 @@ def _tokenize_text(text: str) -> list[str]:
     else:
         bleu = 0.0
 
-    return {
-        "bleu": bleu * 100,  # Return as percentage
-        "precisions": precisions,
-        "bp": bp,
-        "length_ratio": total_hyp_length / total_ref_length
+    return BleuScore(
+        bleu=bleu * 100,  # Return as percentage
+        precisions=precisions,
+        bp=bp,
+        length_ratio=total_hyp_length / total_ref_length
         if total_ref_length > 0
         else 0.0,
-        "hyp_length": total_hyp_length,
-        "ref_length": total_ref_length,
-    }
+        hyp_length=total_hyp_length,
+        ref_length=total_ref_length,
+    )
 
 
 def rouge_score(

pythainlp/tag/__init__.py

@@ -13,9 +13,10 @@
 """
 
 __all__: list[str] = [
-    "PerceptronTagger",
+    "EntitySpan",
     "NER",
     "NNER",
+    "PerceptronTagger",
     "chunk_parse",
     "pos_tag",
     "pos_tag_sents",
@@ -26,5 +27,5 @@
 from pythainlp.tag._tag_perceptron import PerceptronTagger
 from pythainlp.tag.chunk import chunk_parse
 from pythainlp.tag.locations import tag_provinces
-from pythainlp.tag.named_entity import NER, NNER
+from pythainlp.tag.named_entity import NER, NNER, EntitySpan
 from pythainlp.tag.pos_tag import pos_tag, pos_tag_sents, pos_tag_transformers

tests/core/test_tag.py

@@ -8,6 +8,7 @@
 from pythainlp.corpus import download
 from pythainlp.tag import (
     NER,
+    EntitySpan,
     PerceptronTagger,
     perceptron,
     pos_tag,
@@ -318,9 +319,9 @@ def test_get_top_level_entities(self):
 
         # Test with nested entities
         entities = [
-            {"text": ["ห้า"], "span": [7, 9], "entity_type": "cardinal"},
-            {"text": ["ห้า", "โมง"], "span": [7, 11], "entity_type": "time"},
-            {"text": ["โมง"], "span": [9, 11], "entity_type": "unit"},
+            EntitySpan(text=["ห้า"], span=[7, 9], entity_type="cardinal"),
+            EntitySpan(text=["ห้า", "โมง"], span=[7, 11], entity_type="time"),
+            EntitySpan(text=["โมง"], span=[9, 11], entity_type="unit"),
         ]
         top_entities = get_top_level_entities(entities)
         # Should only return 'time' as it contains the others
@@ -330,8 +331,8 @@ def test_get_top_level_entities(self):
 
         # Test with non-overlapping entities
         entities = [
-            {"text": ["วัน"], "span": [0, 1], "entity_type": "time"},
-            {"text": ["เดือน"], "span": [2, 3], "entity_type": "time"},
+            EntitySpan(text=["วัน"], span=[0, 1], entity_type="time"),
+            EntitySpan(text=["เดือน"], span=[2, 3], entity_type="time"),
         ]
         top_entities = get_top_level_entities(entities)
         # Both should be returned as neither contains the other
@@ -352,11 +353,7 @@ def test_entities_to_iob(self):
         # Test basic IOB conversion
         tokens = ["วัน", "ที่", " ", "5", " ", "เมษายน"]
         entities = [
-            {
-                "text": ["5", " ", "เมษายน"],
-                "span": [3, 6],
-                "entity_type": "date",
-            }
+            EntitySpan(text=["5", " ", "เมษายน"], span=[3, 6], entity_type="date")
         ]
         result = _entities_to_iob(tokens, entities)
 
@@ -375,11 +372,7 @@ def test_entities_to_html(self):
         # Test basic HTML conversion
         tokens = ["วัน", "ที่", " ", "5", " ", "เมษายน"]
         entities = [
-            {
-                "text": ["5", " ", "เมษายน"],
-                "span": [3, 6],
-                "entity_type": "date",
-            }
+            EntitySpan(text=["5", " ", "เมษายน"], span=[3, 6], entity_type="date")
         ]
         result = _entities_to_html(tokens, entities)
 
@@ -390,12 +383,8 @@ def test_entities_to_html(self):
         # Test with multiple entities
         tokens = ["นาย", "สมชาย", " ", "อยู่", "ที่", "กรุงเทพ"]
         entities = [
-            {
-                "text": ["นาย", "สมชาย"],
-                "span": [0, 2],
-                "entity_type": "person",
-            },
-            {"text": ["กรุงเทพ"], "span": [5, 6], "entity_type": "location"},
+            EntitySpan(text=["นาย", "สมชาย"], span=[0, 2], entity_type="person"),
+            EntitySpan(text=["กรุงเทพ"], span=[5, 6], entity_type="location"),
         ]
         result = _entities_to_html(tokens, entities)
         expected = "<PERSON>นายสมชาย</PERSON> อยู่ที่<LOCATION>กรุงเทพ</LOCATION>"

tests/extra/testx_benchmarks.py

@@ -7,7 +7,12 @@
 import numpy as np
 import yaml
 
-from pythainlp.benchmarks import bleu_score, rouge_score, word_tokenization
+from pythainlp.benchmarks import (
+    BleuScore,
+    bleu_score,
+    rouge_score,
+    word_tokenization,
+)
 
 with open("./tests/data/sentences.yml", "r", encoding="utf8") as stream:
     TEST_DATA = yaml.safe_load(stream)
@@ -103,6 +108,22 @@ def test_bleu_score_single_reference(self):
         self.assertGreater(score["bleu"], 0)
         self.assertLessEqual(score["bleu"], 100)
 
+    def test_bleu_score_return_type(self):
+        """Test that bleu_score returns a BleuScore typed dict."""
+        references = ["สวัสดีครับ วันนี้อากาศดีมาก"]
+        hypotheses = ["สวัสดีค่ะ วันนี้อากาศดี"]
+
+        score: BleuScore = bleu_score(references, hypotheses)
+
+        self.assertIsInstance(score, dict)
+        self.assertIsInstance(score["bleu"], float)
+        self.assertIsInstance(score["precisions"], list)
+        self.assertTrue(all(isinstance(p, float) for p in score["precisions"]))
+        self.assertIsInstance(score["bp"], float)
+        self.assertIsInstance(score["length_ratio"], float)
+        self.assertIsInstance(score["hyp_length"], int)
+        self.assertIsInstance(score["ref_length"], int)
+
     def test_bleu_score_multiple_references(self):
         """Test BLEU score with multiple references per hypothesis."""
         references = [