diff --git a/pythainlp/benchmarks/__init__.py b/pythainlp/benchmarks/__init__.py
index d55ac6de9..716230cef 100644
--- a/pythainlp/benchmarks/__init__.py
+++ b/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,
diff --git a/pythainlp/benchmarks/metrics.py b/pythainlp/benchmarks/metrics.py
index 8326694ea..b2eb963ee 100644
--- a/pythainlp/benchmarks/metrics.py
+++ b/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(
diff --git a/pythainlp/tag/__init__.py b/pythainlp/tag/__init__.py
index abf4dbdc7..28cb578c5 100644
--- a/pythainlp/tag/__init__.py
+++ b/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
diff --git a/tests/core/test_tag.py b/tests/core/test_tag.py
index 639350f3f..75064afa5 100644
--- a/tests/core/test_tag.py
+++ b/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 = "นายสมชาย อยู่ที่กรุงเทพ"
diff --git a/tests/extra/testx_benchmarks.py b/tests/extra/testx_benchmarks.py
index cdd666cac..5c53d5650 100644
--- a/tests/extra/testx_benchmarks.py
+++ b/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 = [