From fc118dc5a8385163d15fd49e2615bed32362ffaf Mon Sep 17 00:00:00 2001 From: Arthit Suriyawongkul Date: Tue, 24 Mar 2026 07:58:52 +0000 Subject: [PATCH] Use BleuScore TypeDict for BLEU score --- pythainlp/benchmarks/__init__.py | 2 ++ pythainlp/benchmarks/metrics.py | 42 ++++++++++++++++++++------------ pythainlp/tag/__init__.py | 5 ++-- tests/core/test_tag.py | 31 ++++++++--------------- tests/extra/testx_benchmarks.py | 23 ++++++++++++++++- 5 files changed, 64 insertions(+), 39 deletions(-) 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 = [