Repository: Lightning-AI/torchmetrics

Issue:
More tokenizers for the SacreBLEU metric
## 🚀 Feature

Torchmetrics's SacreBLEU is missing some tokenizers from the [SacreBLEU library](https://github.com/mjpost/sacrebleu/tree/master/sacrebleu/tokenizers).

<br>

**Available tokenizers from SacreBLEU**
```python
>import sacrebleu
>print(sacrebleu.BLEU.TOKENIZERS)
... dict_keys(['none', 'zh', '13a', 'intl', 'char', 'ja-mecab', 'ko-mecab', 'spm', 'flores101', 'flores200'])
```
*Just a little note:* `spm` and `flores101` are actually the same tokenizer with the former being under possible removal in the future.

<br>

**Available tokenizers from torchmetrics**
```python
>from torchmetrics.text.sacrebleu import AVAILABLE_TOKENIZERS
>print(AVAILABLE_TOKENIZERS)
... ('none', '13a', 'zh', 'intl', 'char')
```

### Motivation

I think that the SacreBLEU metric implemented by torchmetrics should be as close as possible to the original implementation. 

Adding the missing tokenizers should not be a huge issue as I have already implemented the `ja-mecab` tokenizer (and by implemented I mean shamelessly copied from the source) on [my fork](https://github.com/RistoAle97/torchmetrics/tree/jamecab_tokenizer) because I needed it for a small project of mine.



Discussion:
cc: @stancld 
Cool, I'll try to add them ad the weekend. :]

Resolving patch:
diff --git a/src/torchmetrics/functional/text/sacre_bleu.py b/src/torchmetrics/functional/text/sacre_bleu.py
--- a/src/torchmetrics/functional/text/sacre_bleu.py
+++ b/src/torchmetrics/functional/text/sacre_bleu.py
@@ -37,18 +37,28 @@
 # MIT License
 # Copyright (c) 2017 - Shujian Huang <huangsj@nju.edu.cn>
 
+import os
 import re
+import tempfile
 from functools import partial
-from typing import ClassVar, Optional, Sequence
+from typing import Any, ClassVar, Dict, Optional, Sequence
 
 import torch
 from torch import Tensor, tensor
 from typing_extensions import Literal
 
 from torchmetrics.functional.text.bleu import _bleu_score_compute, _bleu_score_update
-from torchmetrics.utilities.imports import _REGEX_AVAILABLE
+from torchmetrics.utilities.imports import (
+    _IPADIC_AVAILABLE,
+    _MECAB_AVAILABLE,
+    _MECAB_KO_AVAILABLE,
+    _MECAB_KO_DIC_AVAILABLE,
+    _REGEX_AVAILABLE,
+    _SENTENCEPIECE_AVAILABLE,
+)
 
-AVAILABLE_TOKENIZERS = ("none", "13a", "zh", "intl", "char")
+AVAILABLE_TOKENIZERS = ("none", "13a", "zh", "intl", "char", "ja-mecab", "ko-mecab", "flores101", "flores200")
+_Tokenizers_list = Literal["none", "13a", "zh", "intl", "char", "ja-mecab", "ko-mecab", "flores101", "flores200"]
 
 _UCODE_RANGES = (
     ("\u3400", "\u4db5"),  # CJK Unified Ideographs Extension A, release 3.0
@@ -77,6 +87,14 @@
 )
 
 
+_FLORES_LOCAL_DIR = os.path.join(tempfile.gettempdir(), "torchmetrics-flores")
+# Model paths copied from https://github.com/mjpost/sacrebleu/blob/master/sacrebleu/tokenizers/tokenizer_spm.py.
+_FLORES_MODELS_URL = {
+    "flores101": "https://dl.fbaipublicfiles.com/fairseq/models/flores/sacrebleu_tokenizer_spm.model",
+    "flores200": "https://tinyurl.com/flores200sacrebleuspm",
+}
+
+
 class _SacreBLEUTokenizer:
     """Tokenizer used for SacreBLEU calculation.
 
@@ -116,9 +134,18 @@ class _SacreBLEUTokenizer:
         "zh": "_tokenize_zh",
         "intl": "_tokenize_international",
         "char": "_tokenize_char",
+        "ja-mecab": "_tokenize_ja_mecab",
+        "ko-mecab": "_tokenize_ko_mecab",
+        "flores101": "_tokenize_flores_101",
+        "flores200": "_tokenize_flores_200",
     }
 
-    def __init__(self, tokenize: Literal["none", "13a", "zh", "intl", "char"], lowercase: bool = False) -> None:
+    # Keep it as class variable to avoid initializing over and over again
+    sentencepiece_processors: ClassVar[Dict[str, Optional[Any]]] = {"flores101": None, "flores200": None}
+
+    def __init__(self, tokenize: _Tokenizers_list, lowercase: bool = False) -> None:
+        self._check_tokenizers_validity(tokenize)
+
         self.tokenize_fn = getattr(self, self._TOKENIZE_FN[tokenize])
         self.lowercase = lowercase
 
@@ -127,9 +154,9 @@ def __call__(self, line: str) -> Sequence[str]:
         return self._lower(tokenized_line, self.lowercase).split()
 
     @classmethod
-    def tokenize(
-        cls, line: str, tokenize: Literal["none", "13a", "zh", "intl", "char"], lowercase: bool = False
-    ) -> Sequence[str]:
+    def tokenize(cls, line: str, tokenize: _Tokenizers_list, lowercase: bool = False) -> Sequence[str]:
+        cls._check_tokenizers_validity(tokenize)
+
         tokenize_fn = getattr(cls, cls._TOKENIZE_FN[tokenize])
         tokenized_line = tokenize_fn(line)
         return cls._lower(tokenized_line, lowercase).split()
@@ -274,19 +301,159 @@ def _tokenize_char(cls, line: str) -> str:
         """
         return " ".join(char for char in line)
 
+    @classmethod
+    def _tokenize_ja_mecab(cls, line: str) -> str:
+        """Tokenizes a Japanese string line using MeCab morphological analyzer.
+
+        Args:
+            line: the input string to tokenize.
+
+        Return:
+            The tokenized string.
+
+        """
+        import ipadic
+        import MeCab
+
+        tagger = MeCab.Tagger(ipadic.MECAB_ARGS + " -Owakati")
+
+        line = line.strip()
+        return tagger.parse(line).strip()
+
+    @classmethod
+    def _tokenize_ko_mecab(cls, line: str) -> str:
+        """Tokenizes a Korean string line using MeCab-korean morphological analyzer.
+
+        Args:
+            line: the input string to tokenize.
+
+        Return:
+            The tokenized string.
+
+        """
+        import mecab_ko
+        import mecab_ko_dic
+
+        tagger = mecab_ko.Tagger(mecab_ko_dic.MECAB_ARGS + " -Owakati")
+
+        line = line.strip()
+        return tagger.parse(line).strip()
+
+    @classmethod
+    def _tokenize_flores(cls, line: str, tokenize: Literal["flores101", "flores200"]) -> str:
+        """Tokenizes a string line using sentencepiece tokenizer.
+
+        Args:
+            line: the input string to tokenize.
+            tokenize: Tokenization technique to be used.
+
+        Return:
+            The tokenized string.
+
+        """
+        import sentencepiece
+
+        if cls.sentencepiece_processors[tokenize] is None:
+            cls.sentencepiece_processors[tokenize] = sentencepiece.SentencePieceProcessor()
+
+            file_path = os.path.join(_FLORES_LOCAL_DIR, _FLORES_MODELS_URL[tokenize].split("/")[-1])
+            if not os.path.exists(file_path):
+                cls.download_flores_file(tokenize)
+
+            cls.sentencepiece_processors[tokenize].Load(file_path)  # type: ignore[union-attr]
+
+        return " ".join(cls.sentencepiece_processors[tokenize].EncodeAsPieces(line))  # type: ignore[union-attr]
+
+    @classmethod
+    def _tokenize_flores_101(cls, line: str) -> str:
+        """Tokenizes a string line using sentencepiece tokenizer according to `FLORES-101`_ dataset.
+
+        Args:
+            line: the input string to tokenize.
+
+        Return:
+            The tokenized string.
+
+        """
+        return cls._tokenize_flores(line, "flores101")
+
+    @classmethod
+    def _tokenize_flores_200(cls, line: str) -> str:
+        """Tokenizes a string line using sentencepiece tokenizer according to `FLORES-200`_ dataset.
+
+        Args:
+            line: the input string to tokenize.
+
+        Return:
+            The tokenized string.
+
+        """
+        return cls._tokenize_flores(line, "flores200")
+
     @staticmethod
     def _lower(line: str, lowercase: bool) -> str:
         if lowercase:
             return line.lower()
         return line
 
+    @classmethod
+    def _check_tokenizers_validity(cls, tokenize: _Tokenizers_list) -> None:
+        """Check if a supported tokenizer is chosen.
+
+        Also check all dependencies of a given tokenizers are installed.
+
+        """
+        if tokenize not in cls._TOKENIZE_FN:
+            raise ValueError(f"Unsupported tokenizer selected. Please, choose one of {list(cls._TOKENIZE_FN.keys())}")
+
+        if tokenize == "intl" and not _REGEX_AVAILABLE:
+            raise ModuleNotFoundError(
+                "`'intl'` tokenization requires that `regex` is installed."
+                " Use `pip install regex` or `pip install torchmetrics[text]`."
+            )
+
+        if tokenize == "ja-mecab" and not (_MECAB_AVAILABLE and _IPADIC_AVAILABLE):
+            raise ModuleNotFoundError(
+                "`'ja-mecab'` tokenization requires that `MeCab` and `ipadic` are installed."
+                " Use `pip install mecab-python3 ipadic` or `pip install torchmetrics[text]`."
+            )
+
+        if tokenize == "ko-mecab" and not (_MECAB_KO_AVAILABLE and _MECAB_KO_DIC_AVAILABLE):
+            raise ModuleNotFoundError(
+                "`'ko-mecab'` tokenization requires that `mecab_ko` and `mecab_ko_dic` are installed."
+                " Use `pip install mecab_ko mecab_ko_dic` or `pip install torchmetrics[text]`."
+            )
+
+        if "flores" in tokenize and not _SENTENCEPIECE_AVAILABLE:
+            raise ModuleNotFoundError(
+                "`'flores101' and 'flores200'` tokenizations require that `sentencepiece` is installed."
+                " Use `pip install sentencepiece` or `pip install torchmetrics[text]`."
+            )
+
+    @staticmethod
+    def download_flores_file(model_name: Literal["flores101", "flores200"]) -> None:
+        """Download necessary files for `flores` tokenization via `sentencepiece`."""
+        import ssl
+        import urllib.request
+
+        os.makedirs(_FLORES_LOCAL_DIR, exist_ok=True)
+
+        model_url = _FLORES_MODELS_URL[model_name]
+        file_path = os.path.join(_FLORES_LOCAL_DIR, model_url.split("/")[-1])
+
+        try:
+            with open(file_path, "wb") as out_file, urllib.request.urlopen(model_url) as remote_file:
+                out_file.write(remote_file.read())
+        except ssl.SSLError as e:
+            raise OSError(f"Failed to download {model_name} model.") from e
+
 
 def sacre_bleu_score(
     preds: Sequence[str],
     target: Sequence[Sequence[str]],
     n_gram: int = 4,
     smooth: bool = False,
-    tokenize: Literal["none", "13a", "zh", "intl", "char"] = "13a",
+    tokenize: _Tokenizers_list = "13a",
     lowercase: bool = False,
     weights: Optional[Sequence[float]] = None,
 ) -> Tensor:
@@ -299,8 +466,8 @@ def sacre_bleu_score(
         target: An iterable of iterables of reference corpus
         n_gram: Gram value ranged from 1 to 4
         smooth: Whether to apply smoothing - see [2]
-        tokenize: Tokenization technique to be used.
-            Supported tokenization: ['none', '13a', 'zh', 'intl', 'char']
+        tokenize: Tokenization technique to be used. Choose between ``'none'``, ``'13a'``, ``'zh'``, ``'intl'``,
+            ``'char'``, ``'ja-mecab'``, ``'ko-mecab'``, ``'flores101'`` and ``'flores200'``.
         lowercase: If ``True``, BLEU score over lowercased text is calculated.
         weights:
             Weights used for unigrams, bigrams, etc. to calculate BLEU score.
@@ -330,20 +497,8 @@ def sacre_bleu_score(
         and Skip-Bigram Statistics by Chin-Yew Lin and Franz Josef Och `Machine Translation Evolution`_
 
     """
-    if tokenize not in AVAILABLE_TOKENIZERS:
-        raise ValueError(f"Argument `tokenize` expected to be one of {AVAILABLE_TOKENIZERS} but got {tokenize}.")
-
-    if tokenize not in _SacreBLEUTokenizer._TOKENIZE_FN:
-        raise ValueError(
-            f"Unsupported tokenizer selected. Please, choose one of {list(_SacreBLEUTokenizer._TOKENIZE_FN.keys())}"
-        )
     if len(preds) != len(target):
         raise ValueError(f"Corpus has different size {len(preds)} != {len(target)}")
-    if tokenize == "intl" and not _REGEX_AVAILABLE:
-        raise ModuleNotFoundError(
-            "`'intl'` tokenization requires that `regex` is installed."
-            " Use `pip install regex` or `pip install torchmetrics[text]`."
-        )
 
     if weights is not None and len(weights) != n_gram:
         raise ValueError(f"List of weights has different weights than `n_gram`: {len(weights)} != {n_gram}")
diff --git a/src/torchmetrics/text/sacre_bleu.py b/src/torchmetrics/text/sacre_bleu.py
--- a/src/torchmetrics/text/sacre_bleu.py
+++ b/src/torchmetrics/text/sacre_bleu.py
@@ -20,21 +20,17 @@
 from typing import Any, Optional, Sequence, Union
 
 from torch import Tensor
-from typing_extensions import Literal
 
 from torchmetrics.functional.text.bleu import _bleu_score_update
-from torchmetrics.functional.text.sacre_bleu import _SacreBLEUTokenizer
+from torchmetrics.functional.text.sacre_bleu import _SacreBLEUTokenizer, _Tokenizers_list
 from torchmetrics.text.bleu import BLEUScore
-from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE, _REGEX_AVAILABLE
+from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
 
 if not _MATPLOTLIB_AVAILABLE:
     __doctest_skip__ = ["SacreBLEUScore.plot"]
 
 
-AVAILABLE_TOKENIZERS = ("none", "13a", "zh", "intl", "char")
-
-
 class SacreBLEUScore(BLEUScore):
     """Calculate `BLEU score`_ of machine translated text with one or more references.
 
@@ -53,8 +49,8 @@ class SacreBLEUScore(BLEUScore):
     Args:
         n_gram: Gram value ranged from 1 to 4
         smooth: Whether to apply smoothing, see `SacreBLEU`_
-        tokenize: Tokenization technique to be used.
-            Supported tokenization: ``['none', '13a', 'zh', 'intl', 'char']``
+        tokenize: Tokenization technique to be used. Choose between ``'none'``, ``'13a'``, ``'zh'``, ``'intl'``,
+            ``'char'``, ``'ja-mecab'``, ``'ko-mecab'``, ``'flores101'`` and ``'flores200'``.
         lowercase:  If ``True``, BLEU score over lowercased text is calculated.
         kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.
         weights:
@@ -95,20 +91,12 @@ def __init__(
         self,
         n_gram: int = 4,
         smooth: bool = False,
-        tokenize: Literal["none", "13a", "zh", "intl", "char"] = "13a",
+        tokenize: _Tokenizers_list = "13a",
         lowercase: bool = False,
         weights: Optional[Sequence[float]] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(n_gram=n_gram, smooth=smooth, weights=weights, **kwargs)
-        if tokenize not in AVAILABLE_TOKENIZERS:
-            raise ValueError(f"Argument `tokenize` expected to be one of {AVAILABLE_TOKENIZERS} but got {tokenize}.")
-
-        if tokenize == "intl" and not _REGEX_AVAILABLE:
-            raise ModuleNotFoundError(
-                "`'intl'` tokenization requires that `regex` is installed."
-                " Use `pip install regex` or `pip install torchmetrics[text]`."
-            )
         self.tokenizer = _SacreBLEUTokenizer(tokenize, lowercase)
 
     def update(self, preds: Sequence[str], target: Sequence[Sequence[str]]) -> None:
diff --git a/src/torchmetrics/utilities/imports.py b/src/torchmetrics/utilities/imports.py
--- a/src/torchmetrics/utilities/imports.py
+++ b/src/torchmetrics/utilities/imports.py
@@ -58,5 +58,10 @@
 _XLA_AVAILABLE: bool = package_available("torch_xla")
 _PIQ_GREATER_EQUAL_0_8: Optional[bool] = compare_version("piq", operator.ge, "0.8.0")
 _FASTER_COCO_EVAL_AVAILABLE: bool = package_available("faster_coco_eval")
+_MECAB_AVAILABLE: bool = package_available("MeCab")
+_MECAB_KO_AVAILABLE: bool = package_available("mecab_ko")
+_MECAB_KO_DIC_AVAILABLE: bool = package_available("mecab_ko_dic")
+_IPADIC_AVAILABLE: bool = package_available("ipadic")
+_SENTENCEPIECE_AVAILABLE: bool = package_available("sentencepiece")
 
 _LATEX_AVAILABLE: bool = shutil.which("latex") is not None


Test patch:
diff --git a/requirements/text_test.txt b/requirements/text_test.txt
--- a/requirements/text_test.txt
+++ b/requirements/text_test.txt
@@ -5,4 +5,4 @@ jiwer >=2.3.0, <3.1.0
 rouge-score >0.1.0, <=0.1.2
 bert_score ==0.3.13
 huggingface-hub <0.18  # hotfix, failing SDR for latest PT 1.11
-sacrebleu >=2.0.0, <=2.3.1
+sacrebleu >=2.3.0, <=2.3.1
diff --git a/tests/unittests/text/test_sacre_bleu.py b/tests/unittests/text/test_sacre_bleu.py
--- a/tests/unittests/text/test_sacre_bleu.py
+++ b/tests/unittests/text/test_sacre_bleu.py
@@ -17,7 +17,7 @@
 
 import pytest
 from torch import Tensor, tensor
-from torchmetrics.functional.text.sacre_bleu import sacre_bleu_score
+from torchmetrics.functional.text.sacre_bleu import AVAILABLE_TOKENIZERS, _Tokenizers_list, sacre_bleu_score
 from torchmetrics.text.sacre_bleu import SacreBLEUScore
 from torchmetrics.utilities.imports import _SACREBLEU_AVAILABLE
 
@@ -28,9 +28,6 @@
     from sacrebleu.metrics import BLEU
 
 
-TOKENIZERS = ("none", "13a", "zh", "intl", "char")
-
-
 def _sacrebleu_fn(preds: Sequence[str], targets: Sequence[Sequence[str]], tokenize: str, lowercase: bool) -> Tensor:
     sacrebleu_fn = BLEU(tokenize=tokenize, lowercase=lowercase)
     # Sacrebleu expects different format of input
@@ -44,7 +41,7 @@ def _sacrebleu_fn(preds: Sequence[str], targets: Sequence[Sequence[str]], tokeni
     [(_inputs_multiple_references.preds, _inputs_multiple_references.targets)],
 )
 @pytest.mark.parametrize(["lowercase"], [(False,), (True,)])
-@pytest.mark.parametrize("tokenize", TOKENIZERS)
+@pytest.mark.parametrize("tokenize", AVAILABLE_TOKENIZERS)
 @pytest.mark.skipif(not _SACREBLEU_AVAILABLE, reason="test requires sacrebleu")
 class TestSacreBLEUScore(TextTester):
     """Test class for `SacreBLEUScore` metric."""
@@ -109,3 +106,26 @@ def test_no_and_uniform_weights_class():
     no_weights_score = no_weights_bleu(preds, targets)
     uniform_weights_score = uniform_weights_bleu(preds, targets)
     assert no_weights_score == uniform_weights_score
+
+
+def test_tokenize_ja_mecab():
+    """Test that `ja-mecab` tokenizer works on a Japanese text in alignment with the SacreBleu implementation."""
+    sacrebleu = SacreBLEUScore(tokenize="ja-mecab")
+
+    preds = ["これは美しい花です。"]
+    targets = [["これは美しい花です。", "おいしい寿司を食べたい。"]]
+    assert sacrebleu(preds, targets) == _sacrebleu_fn(preds, targets, tokenize="ja-mecab", lowercase=False)
+
+
+def test_tokenize_ko_mecab():
+    """Test that `ja-mecab` tokenizer works on a Japanese text in alignment with the SacreBleu implementation."""
+    sacrebleu = SacreBLEUScore(tokenize="ko-mecab")
+
+    preds = ["이 책은 정말 재미있어요."]
+    targets = [["이 책은 정말 재미있어요.", "고마워요, 너무 도와줘서."]]
+    assert sacrebleu(preds, targets) == _sacrebleu_fn(preds, targets, tokenize="ko-mecab", lowercase=False)
+
+
+def test_equivalence_of_available_tokenizers_and_annotation():
+    """Test equivalence of SacreBLEU available tokenizers and corresponding type annotation."""
+    assert set(AVAILABLE_TOKENIZERS) == set(_Tokenizers_list.__args__)
diff --git a/tests/unittests/utilities/test_plot.py b/tests/unittests/utilities/test_plot.py
--- a/tests/unittests/utilities/test_plot.py
+++ b/tests/unittests/utilities/test_plot.py
@@ -11,6 +11,7 @@
 # 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.
+import sys
 from functools import partial
 from typing import Callable
 
@@ -721,6 +722,7 @@ def test_plot_methods_special_image_metrics(metric_class, preds, target, index_0
     plt.close(fig)
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="DDP not supported on windows")
 def test_plot_methods_special_text_metrics():
     """Test the plot method for text metrics that does not fit the default testing format."""
     metric = BERTScore()