sacre_bleu.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555
  1. # Copyright The Lightning team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. # referenced from
  15. # Library Name: torchtext
  16. # Authors: torchtext authors and @sluks
  17. # Date: 2020-07-18
  18. # Link: https://pytorch.org/text/_modules/torchtext/data/metrics.html#bleu_score
  19. ##############
  20. # Copyright 2017--2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  21. #
  22. # Licensed under the Apache License, Version 2.0 (the "License"). You may not
  23. # use this file except in compliance with the License. A copy of the License
  24. # is located at
  25. #
  26. # http://aws.amazon.com/apache2.0/
  27. #
  28. # or in the "license" file accompanying this file. This file is distributed on
  29. # an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
  30. # express or implied. See the License for the specific language governing
  31. # permissions and limitations under the License.
  32. ##############
  33. # MIT License
  34. # Copyright (c) 2017 - Shujian Huang <huangsj@nju.edu.cn>
  35. import os
  36. import re
  37. import tempfile
  38. from collections.abc import Sequence
  39. from functools import partial
  40. from typing import Any, ClassVar, Optional
  41. import torch
  42. from torch import Tensor, tensor
  43. from typing_extensions import Literal
  44. from torchmetrics.functional.text.bleu import _bleu_score_compute, _bleu_score_update
  45. from torchmetrics.utilities.imports import (
  46. _IPADIC_AVAILABLE,
  47. _MECAB_AVAILABLE,
  48. _MECAB_KO_AVAILABLE,
  49. _MECAB_KO_DIC_AVAILABLE,
  50. _REGEX_AVAILABLE,
  51. _SENTENCEPIECE_AVAILABLE,
  52. )
  53. AVAILABLE_TOKENIZERS = ("none", "13a", "zh", "intl", "char", "ja-mecab", "ko-mecab", "flores101", "flores200")
  54. _TokenizersLiteral = Literal["none", "13a", "zh", "intl", "char", "ja-mecab", "ko-mecab", "flores101", "flores200"]
  55. _UCODE_RANGES = (
  56. ("\u3400", "\u4db5"), # CJK Unified Ideographs Extension A, release 3.0
  57. ("\u4e00", "\u9fa5"), # CJK Unified Ideographs, release 1.1
  58. ("\u9fa6", "\u9fbb"), # CJK Unified Ideographs, release 4.1
  59. ("\uf900", "\ufa2d"), # CJK Compatibility Ideographs, release 1.1
  60. ("\ufa30", "\ufa6a"), # CJK Compatibility Ideographs, release 3.2
  61. ("\ufa70", "\ufad9"), # CJK Compatibility Ideographs, release 4.1
  62. ("\u20000", "\u2a6d6"), # (UTF16) CJK Unified Ideographs Extension B, release 3.1
  63. ("\u2f800", "\u2fa1d"), # (UTF16) CJK Compatibility Supplement, release 3.1
  64. ("\uff00", "\uffef"), # Full width ASCII, full width of English punctuation,
  65. # half width Katakana, half wide half width kana, Korean alphabet
  66. ("\u2e80", "\u2eff"), # CJK Radicals Supplement
  67. ("\u3000", "\u303f"), # CJK punctuation mark
  68. ("\u31c0", "\u31ef"), # CJK stroke
  69. ("\u2f00", "\u2fdf"), # Kangxi Radicals
  70. ("\u2ff0", "\u2fff"), # Chinese character structure
  71. ("\u3100", "\u312f"), # Phonetic symbols
  72. ("\u31a0", "\u31bf"), # Phonetic symbols (Taiwanese and Hakka expansion)
  73. ("\ufe10", "\ufe1f"),
  74. ("\ufe30", "\ufe4f"),
  75. ("\u2600", "\u26ff"),
  76. ("\u2700", "\u27bf"),
  77. ("\u3200", "\u32ff"),
  78. ("\u3300", "\u33ff"),
  79. )
  80. _FLORES_LOCAL_DIR = os.path.join(tempfile.gettempdir(), "torchmetrics-flores")
  81. # Model paths copied from https://github.com/mjpost/sacrebleu/blob/master/sacrebleu/tokenizers/tokenizer_spm.py.
  82. _FLORES_MODELS_URL = {
  83. "flores101": "https://dl.fbaipublicfiles.com/fairseq/models/flores/sacrebleu_tokenizer_spm.model",
  84. "flores200": "https://tinyurl.com/flores200sacrebleuspm",
  85. }
  86. class _SacreBLEUTokenizer:
  87. """Tokenizer used for SacreBLEU calculation.
  88. Source: https://github.com/mjpost/sacrebleu/tree/master/sacrebleu/tokenizers
  89. """
  90. _REGEX = (
  91. # language-dependent part (assuming Western languages)
  92. (re.compile(r"([\{-\~\[-\` -\&\(-\+\:-\@\/])"), r" \1 "),
  93. # tokenize period and comma unless preceded by a digit
  94. (re.compile(r"([^0-9])([\.,])"), r"\1 \2 "),
  95. # tokenize period and comma unless followed by a digit
  96. (re.compile(r"([\.,])([^0-9])"), r" \1 \2"),
  97. # tokenize dash when preceded by a digit
  98. (re.compile(r"([0-9])(-)"), r"\1 \2 "),
  99. # one space only between words
  100. # NOTE: Doing this in Python (below) is faster
  101. # (re.compile(r'\s+'), r' '),
  102. )
  103. if _REGEX_AVAILABLE:
  104. import regex
  105. _INT_REGEX = (
  106. # Separate out punctuation preceded by a non-digit
  107. (regex.compile(r"(\P{N})(\p{P})"), r"\1 \2 "),
  108. # Separate out punctuation followed by a non-digit
  109. (regex.compile(r"(\p{P})(\P{N})"), r" \1 \2"),
  110. # Separate out symbols
  111. (regex.compile(r"(\p{S})"), r" \1 "),
  112. )
  113. _TOKENIZE_FN: ClassVar[dict] = {
  114. "none": "_tokenize_base",
  115. "13a": "_tokenize_13a",
  116. "zh": "_tokenize_zh",
  117. "intl": "_tokenize_international",
  118. "char": "_tokenize_char",
  119. "ja-mecab": "_tokenize_ja_mecab",
  120. "ko-mecab": "_tokenize_ko_mecab",
  121. "flores101": "_tokenize_flores_101",
  122. "flores200": "_tokenize_flores_200",
  123. }
  124. # Keep it as class variable to avoid initializing over and over again
  125. sentencepiece_processors: ClassVar[dict[str, Optional[Any]]] = {"flores101": None, "flores200": None}
  126. def __init__(self, tokenize: _TokenizersLiteral, lowercase: bool = False) -> None:
  127. self._check_tokenizers_validity(tokenize)
  128. self.tokenize_fn = getattr(self, self._TOKENIZE_FN[tokenize])
  129. self.lowercase = lowercase
  130. def __call__(self, line: str) -> Sequence[str]:
  131. tokenized_line = self.tokenize_fn(line)
  132. return self._lower(tokenized_line, self.lowercase).split()
  133. @classmethod
  134. def tokenize(
  135. cls: type["_SacreBLEUTokenizer"],
  136. line: str,
  137. tokenize: _TokenizersLiteral,
  138. lowercase: bool = False,
  139. ) -> Sequence[str]:
  140. cls._check_tokenizers_validity(tokenize)
  141. tokenize_fn = getattr(cls, cls._TOKENIZE_FN[tokenize])
  142. tokenized_line = tokenize_fn(line)
  143. return cls._lower(tokenized_line, lowercase).split()
  144. @classmethod
  145. def _tokenize_regex(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  146. """Post-processing tokenizer for `13a` and `zh` tokenizers.
  147. Args:
  148. line: a segment to tokenize
  149. Return:
  150. the tokenized line
  151. """
  152. for _re, repl in cls._REGEX:
  153. line = _re.sub(repl, line)
  154. # no leading or trailing spaces, single space within words
  155. return " ".join(line.split())
  156. @staticmethod
  157. def _is_chinese_char(uchar: str) -> bool:
  158. """Check if character is chinese.
  159. Args:
  160. uchar: input char in unicode.
  161. Return:
  162. whether the input char is a Chinese character.
  163. """
  164. return any(start <= uchar <= end for start, end in _UCODE_RANGES)
  165. @classmethod
  166. def _tokenize_base(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  167. """Tokenizes an input line with the tokenizer.
  168. Args:
  169. line: a segment to tokenize
  170. Return:
  171. the tokenized line
  172. """
  173. return line
  174. @classmethod
  175. def _tokenize_13a(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  176. """Tokenizes a line using a relatively minimal tokenization that is equivalent to mteval-v13a, used by WMT.
  177. Args:
  178. line: input sentence
  179. Return:
  180. tokenized sentence
  181. """
  182. # language-independent part:
  183. line = line.replace("<skipped>", "")
  184. line = line.replace("-\n", "")
  185. line = line.replace("\n", " ")
  186. if "&" in line:
  187. line = line.replace("&quot;", '"')
  188. line = line.replace("&amp;", "&")
  189. line = line.replace("&lt;", "<")
  190. line = line.replace("&gt;", ">")
  191. return cls._tokenize_regex(f" {line} ")
  192. @classmethod
  193. def _tokenize_zh(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  194. """Tokenization of Chinese text.
  195. This is done in two steps: separate each Chinese characters (by utf-8 encoding) and afterwards tokenize the
  196. Chinese part (following the `13a` i.e. mteval tokenizer).
  197. Author: Shujian Huang huangsj@nju.edu.cn.
  198. Args:
  199. line: input sentence
  200. Return:
  201. tokenized sentence
  202. """
  203. line = line.strip()
  204. line_in_chars = ""
  205. for char in line:
  206. if cls._is_chinese_char(char):
  207. line_in_chars += " "
  208. line_in_chars += char
  209. line_in_chars += " "
  210. else:
  211. line_in_chars += char
  212. return cls._tokenize_regex(line_in_chars)
  213. @classmethod
  214. def _tokenize_international(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  215. r"""Tokenizes a string following the official BLEU implementation.
  216. See github.com/moses-smt/mosesdecoder/blob/master/scripts/generic/mteval-v14.pl#L954-L983
  217. In our case, the input string is expected to be just one line.
  218. We just tokenize on punctuation and symbols,
  219. except when a punctuation is preceded and followed by a digit
  220. (e.g. a comma/dot as a thousand/decimal separator).
  221. We do not recover escaped forms of punctuation such as &apos; or &gt;
  222. as these should never appear in MT system outputs (see issue #138)
  223. Note that a number (e.g., a year) followed by a dot at the end of
  224. sentence is NOT tokenized, i.e. the dot stays with the number because
  225. `s/(\\p{P})(\\P{N})/ $1 $2/g` does not match this case (unless we add a
  226. space after each sentence). However, this error is already in the
  227. original mteval-v14.pl and we want to be consistent with it.
  228. The error is not present in the non-international version,
  229. which uses `$norm_text = " $norm_text "`.
  230. Args:
  231. line: the input string to tokenize.
  232. Return:
  233. The tokenized string.
  234. """
  235. for _re, repl in cls._INT_REGEX:
  236. line = _re.sub(repl, line)
  237. return " ".join(line.split())
  238. @classmethod
  239. def _tokenize_char(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  240. """Tokenizes all the characters in the input line.
  241. Args:
  242. line: a segment to tokenize
  243. Return:
  244. the tokenized line
  245. """
  246. return " ".join(char for char in line)
  247. @classmethod
  248. def _tokenize_ja_mecab(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  249. """Tokenizes a Japanese string line using MeCab morphological analyzer.
  250. Args:
  251. line: the input string to tokenize.
  252. Return:
  253. The tokenized string.
  254. """
  255. import ipadic
  256. import MeCab
  257. tagger = MeCab.Tagger(ipadic.MECAB_ARGS + " -Owakati")
  258. line = line.strip()
  259. return tagger.parse(line).strip()
  260. @classmethod
  261. def _tokenize_ko_mecab(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  262. """Tokenizes a Korean string line using MeCab-korean morphological analyzer.
  263. Args:
  264. line: the input string to tokenize.
  265. Return:
  266. The tokenized string.
  267. """
  268. import mecab_ko
  269. import mecab_ko_dic
  270. tagger = mecab_ko.Tagger(mecab_ko_dic.MECAB_ARGS + " -Owakati")
  271. line = line.strip()
  272. return tagger.parse(line).strip()
  273. @classmethod
  274. def _tokenize_flores(
  275. cls: type["_SacreBLEUTokenizer"], line: str, tokenize: Literal["flores101", "flores200"]
  276. ) -> str:
  277. """Tokenizes a string line using sentencepiece tokenizer.
  278. Args:
  279. line: the input string to tokenize.
  280. tokenize: Tokenization technique to be used.
  281. Return:
  282. The tokenized string.
  283. """
  284. import sentencepiece
  285. if cls.sentencepiece_processors[tokenize] is None:
  286. cls.sentencepiece_processors[tokenize] = sentencepiece.SentencePieceProcessor()
  287. file_path = os.path.join(_FLORES_LOCAL_DIR, _FLORES_MODELS_URL[tokenize].split("/")[-1])
  288. if not os.path.exists(file_path):
  289. cls.download_flores_file(tokenize)
  290. cls.sentencepiece_processors[tokenize].Load(file_path) # type: ignore[union-attr]
  291. return " ".join(cls.sentencepiece_processors[tokenize].EncodeAsPieces(line)) # type: ignore[union-attr]
  292. @classmethod
  293. def _tokenize_flores_101(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  294. """Tokenizes a string line using sentencepiece tokenizer according to `FLORES-101`_ dataset.
  295. Args:
  296. line: the input string to tokenize.
  297. Return:
  298. The tokenized string.
  299. """
  300. return cls._tokenize_flores(line, "flores101")
  301. @classmethod
  302. def _tokenize_flores_200(cls: type["_SacreBLEUTokenizer"], line: str) -> str:
  303. """Tokenizes a string line using sentencepiece tokenizer according to `FLORES-200`_ dataset.
  304. Args:
  305. line: the input string to tokenize.
  306. Return:
  307. The tokenized string.
  308. """
  309. return cls._tokenize_flores(line, "flores200")
  310. @staticmethod
  311. def _lower(line: str, lowercase: bool) -> str:
  312. if lowercase:
  313. return line.lower()
  314. return line
  315. @classmethod
  316. def _check_tokenizers_validity(cls: type["_SacreBLEUTokenizer"], tokenize: _TokenizersLiteral) -> None:
  317. """Check if a supported tokenizer is chosen.
  318. Also check all dependencies of a given tokenizers are installed.
  319. """
  320. if tokenize not in cls._TOKENIZE_FN:
  321. raise ValueError(f"Unsupported tokenizer selected. Please, choose one of {list(cls._TOKENIZE_FN.keys())}")
  322. if tokenize == "intl" and not _REGEX_AVAILABLE:
  323. raise ModuleNotFoundError(
  324. "`'intl'` tokenization requires that `regex` is installed."
  325. " Use `pip install regex` or `pip install torchmetrics[text]`."
  326. )
  327. if tokenize == "ja-mecab" and not (_MECAB_AVAILABLE and _IPADIC_AVAILABLE):
  328. raise ModuleNotFoundError(
  329. "`'ja-mecab'` tokenization requires that `MeCab` and `ipadic` are installed."
  330. " Use `pip install mecab-python3 ipadic` or `pip install torchmetrics[text]`."
  331. )
  332. if tokenize == "ko-mecab" and not (_MECAB_KO_AVAILABLE and _MECAB_KO_DIC_AVAILABLE):
  333. raise ModuleNotFoundError(
  334. "`'ko-mecab'` tokenization requires that `mecab_ko` and `mecab_ko_dic` are installed."
  335. " Use `pip install mecab_ko mecab_ko_dic` or `pip install torchmetrics[text]`."
  336. )
  337. if "flores" in tokenize and not _SENTENCEPIECE_AVAILABLE:
  338. raise ModuleNotFoundError(
  339. "`'flores101' and 'flores200'` tokenizations require that `sentencepiece` is installed."
  340. " Use `pip install sentencepiece` or `pip install torchmetrics[text]`."
  341. )
  342. @staticmethod
  343. def download_flores_file(model_name: Literal["flores101", "flores200"]) -> None:
  344. """Download necessary files for `flores` tokenization via `sentencepiece`."""
  345. import ssl
  346. import urllib.request
  347. os.makedirs(_FLORES_LOCAL_DIR, exist_ok=True)
  348. model_url = _FLORES_MODELS_URL[model_name]
  349. file_path = os.path.join(_FLORES_LOCAL_DIR, model_url.split("/")[-1])
  350. try:
  351. with open(file_path, "wb") as out_file, urllib.request.urlopen(model_url) as remote_file:
  352. out_file.write(remote_file.read())
  353. except ssl.SSLError as e:
  354. raise OSError(f"Failed to download {model_name} model.") from e
  355. def sacre_bleu_score(
  356. preds: Sequence[str],
  357. target: Sequence[Sequence[str]],
  358. n_gram: int = 4,
  359. smooth: bool = False,
  360. tokenize: _TokenizersLiteral = "13a",
  361. lowercase: bool = False,
  362. weights: Optional[Sequence[float]] = None,
  363. ) -> Tensor:
  364. """Calculate `BLEU score`_ [1] of machine translated text with one or more references.
  365. This implementation follows the behaviour of SacreBLEU [2] implementation from https://github.com/mjpost/sacrebleu.
  366. .. note::
  367. In the original SacreBLEU, references are passed as a list of reference sets (grouped by reference index).
  368. In TorchMetrics, references are passed grouped per prediction (each prediction has its own list of references).
  369. For example::
  370. # Predictions
  371. preds = ['The dog bit the man.', "It wasn't surprising.", 'The man had just bitten him.']
  372. # Original SacreBLEU:
  373. refs = [
  374. ['The dog bit the man.', 'It was not unexpected.', 'The man bit him first.'], # First set
  375. ['The dog had bit the man.', 'No one was surprised.', 'The man had bitten the dog.'], # Second set
  376. ]
  377. # TorchMetrics SacreBLEU:
  378. target = [
  379. ['The dog bit the man.', 'The dog had bit the man.'], # References for first prediction
  380. ['It was not unexpected.', 'No one was surprised.'], # References for second prediction
  381. ['The man bit him first.', 'The man had bitten the dog.'], # References for third prediction
  382. ]
  383. Args:
  384. preds: An iterable of machine translated corpus
  385. target: An iterable of iterables of reference corpus
  386. n_gram: Gram value ranged from 1 to 4
  387. smooth: Whether to apply smoothing - see [2]
  388. tokenize: Tokenization technique to be used. Choose between ``'none'``, ``'13a'``, ``'zh'``, ``'intl'``,
  389. ``'char'``, ``'ja-mecab'``, ``'ko-mecab'``, ``'flores101'`` and ``'flores200'``.
  390. lowercase: If ``True``, BLEU score over lowercased text is calculated.
  391. weights:
  392. Weights used for unigrams, bigrams, etc. to calculate BLEU score.
  393. If not provided, uniform weights are used.
  394. Return:
  395. Tensor with BLEU Score
  396. Raises:
  397. ValueError: If ``preds`` and ``target`` corpus have different lengths.
  398. ValueError: If a length of a list of weights is not ``None`` and not equal to ``n_gram``.
  399. Example:
  400. >>> from torchmetrics.functional.text import sacre_bleu_score
  401. >>> preds = ['the cat is on the mat']
  402. >>> target = [['there is a cat on the mat', 'a cat is on the mat']]
  403. >>> sacre_bleu_score(preds, target)
  404. tensor(0.7598)
  405. References:
  406. [1] BLEU: a Method for Automatic Evaluation of Machine Translation by Papineni,
  407. Kishore, Salim Roukos, Todd Ward, and Wei-Jing Zhu `BLEU`_
  408. [2] A Call for Clarity in Reporting BLEU Scores by Matt Post.
  409. [3] Automatic Evaluation of Machine Translation Quality Using Longest Common Subsequence
  410. and Skip-Bigram Statistics by Chin-Yew Lin and Franz Josef Och `Machine Translation Evolution`_
  411. """
  412. if len(preds) != len(target):
  413. raise ValueError(f"Corpus has different size {len(preds)} != {len(target)}")
  414. if weights is not None and len(weights) != n_gram:
  415. raise ValueError(f"List of weights has different weights than `n_gram`: {len(weights)} != {n_gram}")
  416. if weights is None:
  417. weights = [1.0 / n_gram] * n_gram
  418. numerator = torch.zeros(n_gram)
  419. denominator = torch.zeros(n_gram)
  420. preds_len = tensor(0.0)
  421. target_len = tensor(0.0)
  422. tokenize_fn = partial(_SacreBLEUTokenizer.tokenize, tokenize=tokenize, lowercase=lowercase)
  423. preds_len, target_len = _bleu_score_update(
  424. preds,
  425. target,
  426. numerator,
  427. denominator,
  428. preds_len,
  429. target_len,
  430. n_gram,
  431. tokenize_fn,
  432. )
  433. return _bleu_score_compute(preds_len, target_len, numerator, denominator, n_gram, weights, smooth)