chinesebleu.py

# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Chinese BLEU"""

import evaluate
import datasets
import math
from collections import Counter
#import jieba_fast as jieba
import pycantonese

# TODO: Add BibTeX citation
#_CITATION = """\
#@InProceedings{huggingface:module,
#title = {A great new module},
#authors={huggingface, Inc.},
#year={2020}
#}
#"""
_CITATION = ""

# TODO: Add description of the module here
_DESCRIPTION = """\
This evaluation metric is tailor-made to evaluate the translation quality of Chinese translation using customized implementation of BLEU evaluation metric.
"""


# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates how good are predictions given some references, using certain scores
Args:
    predictions (str): translation sentence to score.
    references (str): reference sentence for each translation.
Returns:
    score: the Chinese BLEU score,
    counts: Counts in n-gram (1-4 grams),
    totals: Totals in n-gram,
    bp: Brevity Penalty,
    tokenizer: Selection of Tokenizer (either "chinese" or "char")
    
Examples:
    Examples should be written in doctest format, and should illustrate how
    to use the function.

    >>> my_new_module = evaluate.load("chinesebleu")
    >>> results = my_new_module.compute(references=["這裡就是香港都會大學"], predictions=["這裡是香港都會大學"])
    >>> print(results)
    {'score': 71.89393375176813, 'counts': [9, 7, 5, 4], 'totals': [9, 8, 7, 6], 'bp': 1.0, 'sys_len': 9, 'ref_len': 10, tokenizer: 'chinese'}
"""

@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class ChineseBLEU(evaluate.Metric):
    """Chinese BLEU - a BLEU-based metric for Chinese sentences"""

    def _info(self):
        return evaluate.MetricInfo(
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            features=datasets.Features({
                'predictions': datasets.Value('string'),
                'references': datasets.Value('string'),
            }),
            homepage="https://huggingface.co/spaces/raptorkwok/chinesebleu/",
            codebase_urls=["https://huggingface.co/spaces/raptorkwok/chinesebleu/"]
        )

    def _download_and_prepare(self, dl_manager):
        """No extra files required to download, pass"""
        pass

    def _tokenize_chinese(self, sentence, tokenizer='char'):
        """
        Tokenize Chinese sentence.
        
        Args:
            sentence (str): Input Chinese sentence.
            tokenizer (str): 'char' for character-level, 'chinese' for word-level segmentation.
        
        Returns:
            list: List of tokens.
        """
        if tokenizer == 'chinese':
            #return list(jieba.cut(sentence, cut_all=False))
            return pycantonese.segment(sentence)
        else:
            return list(sentence)  # Character-level tokenization

    def _get_ngrams(self, tokens, n):
        """
        Extract n-grams from a list of tokens.
        
        Args:
            tokens (list): List of tokens.
            n (int): N-gram order.
        
        Returns:
            list: List of n-grams as tuples.
        """
        return [tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)]
        
    def _compute(self, predictions, references):
        """
        Compute BLEU score for a corpus of predictions against references.
        
        Assumes one reference per prediction. For multiple references per prediction,
        modify the clipping logic accordingly.
        
        Args:
            predictions (list[str]): List of predicted sentences.
            references (list[str]): List of reference sentences (same length as predictions).
        
        Returns:
            tuple: (bleu_score, counts, totals, precisions, brevity_penalty)
                - score (float): BLEU score (0-100).
                - counts (list[int]): Clipped n-gram counts [c1, c2, ..., cN].
                - totals (list[int]): Total n-gram counts [t1, t2, ..., tN].
                - precisions (list[float]): N-gram precisions [p1, p2, ..., pN].
                - brevity_penalty (float): Brevity penalty (0-1).
        """
        if len(predictions) != len(references):
            raise ValueError("Predictions and references must have the same length.")

        max_n = 4 # Default n-gram = 4

        counts = [0] * max_n
        totals = [0] * max_n
        
        tokenizer = 'chinese'
        
        pred_tokens = [self._tokenize_chinese(p, tokenizer) for p in predictions]
        ref_tokens = [self._tokenize_chinese(r, tokenizer) for r in references]

        # For total number of tokens < 4, fallback to SacreBLEU
        if len(pred_tokens[0]) < 4 or len(ref_tokens[0]) < 4:
            tokenizer = 'char'
            sacrebleu = evaluate.load('sacrebleu')
            bleu_result = sacrebleu.compute(predictions=predictions, references=references, tokenize="zh")
            bleu_result['tokenizer'] = tokenizer
            return bleu_result
        
        for n in range(1, max_n + 1):
            clipped_counts_n = 0
            total_ngrams_n = 0
            
            for ptoks, rtoks in zip(pred_tokens, ref_tokens):
                pred_ngrams = self._get_ngrams(ptoks, n)
                total_ngrams_n += len(pred_ngrams)
                
                ref_ngrams_count = Counter(self._get_ngrams(rtoks, n))
                pred_count = Counter(pred_ngrams)
                
                for ngram, count in pred_count.items():
                    clipped = min(count, ref_ngrams_count.get(ngram, 0))
                    clipped_counts_n += clipped
            
            counts[n - 1] = clipped_counts_n
            totals[n - 1] = total_ngrams_n
        
        # Compute precisions
        precisions = []
        for c, t in zip(counts, totals):
            if t == 0:
                precisions.append(0.0)
            else:
                precisions.append(float(c) / t)
        
        # Geometric mean of precisions
        if any(p == 0 for p in precisions):
            geom_mean = 0.0
        else:
            log_sum = sum(math.log(p) for p in precisions)
            geom_mean = math.exp(log_sum / max_n)
        
        # Brevity penalty
        c_len = sum(len(pt) for pt in pred_tokens)
        r_len = sum(len(rt) for rt in ref_tokens)
        
        if c_len == 0:
            bp = 0.0
        elif c_len >= r_len:
            bp = 1.0
        else:
            bp = math.exp(1 - r_len / c_len)
        
        bleu = bp * geom_mean * 100
        return {
            "score": bleu, 
            "counts": counts, 
            "totals": totals, 
            "precisions": precisions, 
            "ref_len": len(ref_tokens[0]),
            "sys_len": len(pred_tokens[0]),
            "bp": bp,
            "tokenizer": tokenizer
        }