| import os |
| import re |
| import unicodedata |
| from typing import Dict |
|
|
| import kenlm |
| import sentencepiece |
| from huggingface_hub import cached_download, hf_hub_url |
|
|
|
|
| class SentencePiece: |
| def __init__( |
| self, |
| model: str, |
| ): |
| super().__init__() |
| self.sp = sentencepiece.SentencePieceProcessor() |
| self.sp.load(str(model)) |
|
|
| def do(self, text: dict) -> dict: |
| tokenized = self.sp.encode_as_pieces(text) |
| return " ".join(tokenized) |
|
|
|
|
| class KenlmModel: |
| digit_re: re.Pattern = re.compile(r"\d") |
| unicode_punct: Dict[str, str] = { |
| "，": ",", |
| "。": ".", |
| "、": ",", |
| "„": '"', |
| "”": '"', |
| "“": '"', |
| "«": '"', |
| "»": '"', |
| "１": '"', |
| "」": '"', |
| "「": '"', |
| "《": '"', |
| "》": '"', |
| "´": "'", |
| "∶": ":", |
| "：": ":", |
| "？": "?", |
| "！": "!", |
| "（": "(", |
| "）": ")", |
| "；": ";", |
| "–": "-", |
| "—": " - ", |
| "．": ". ", |
| "～": "~", |
| "’": "'", |
| "…": "...", |
| "━": "-", |
| "〈": "<", |
| "〉": ">", |
| "【": "[", |
| "】": "]", |
| "％": "%", |
| "►": "-", |
| } |
| unicode_punct_re = re.compile(f"[{''.join(unicode_punct.keys())}]") |
| non_printing_chars_re = re.compile( |
| f"[{''.join(map(chr, list(range(0,32)) + list(range(127,160))))}]" |
| ) |
| kenlm_model_dir = None |
| sentence_piece_model_dir = None |
|
|
| def __init__( |
| self, |
| model_dataset: str, |
| language: str, |
| lower_case: bool = False, |
| remove_accents: bool = False, |
| normalize_numbers: bool = True, |
| punctuation: int = 1, |
| ): |
| self.model = kenlm.Model(os.path.join(model_dataset, f"{language}.arpa.bin")) |
| self.tokenizer = SentencePiece(os.path.join(model_dataset, f"{language}.sp.model")) |
| self.accent = remove_accents |
| self.case = lower_case |
| self.numbers = normalize_numbers |
| self.punct = punctuation |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| model_dataset: str, |
| language: str, |
| ): |
| return cls( |
| model_dataset, |
| language, |
| False, |
| False, |
| True, |
| 1, |
| ) |
|
|
| def pp(self, log_score, length): |
| return 10.0 ** (-log_score / length) |
|
|
| def get_perplexity(self, doc: str, normalize_cc_net: bool = True): |
| if normalize_cc_net: |
| doc = self.normalize( |
| doc, |
| accent=self.accent, |
| case=self.case, |
| numbers=self.numbers, |
| punct=self.punct, |
| ) |
| |
| doc = self.tokenizer.do(doc) |
| doc_log_score, doc_length = 0, 0 |
| for line in doc.split("\n"): |
| log_score = self.model.score(line) |
| length = len(line.split()) + 1 |
| doc_log_score += log_score |
| doc_length += length |
| return round(self.pp(doc_log_score, doc_length), 1) |
|
|
| def normalize( |
| self, |
| line: str, |
| accent: bool = True, |
| case: bool = True, |
| numbers: bool = True, |
| punct: int = 1, |
| ) -> str: |
| line = line.strip() |
| if not line: |
| return line |
| if case: |
| line = line.lower() |
| if accent: |
| line = self.strip_accents(line) |
| if numbers: |
| line = self.digit_re.sub("0", line) |
| if punct == 1: |
| line = self.replace_unicode_punct(line) |
| elif punct == 2: |
| line = self.remove_unicode_punct(line) |
| line = self.remove_non_printing_char(line) |
| return line |
|
|
| def strip_accents(self, line: str) -> str: |
| """Strips accents from a piece of text.""" |
| nfd = unicodedata.normalize("NFD", line) |
| output = [c for c in nfd if unicodedata.category(c) != "Mn"] |
| if len(output) == line: |
| return line |
| return "".join(output) |
|
|
| def replace_unicode_punct(self, text: str) -> str: |
| return "".join(self.unicode_punct.get(c, c) for c in text) |
|
|
| def remove_unicode_punct(self, text: str) -> str: |
| """More aggressive version of replace_unicode_punct but also faster.""" |
| return self.unicode_punct_re.sub("", text) |
|
|
| def remove_non_printing_char(self, text: str) -> str: |
| return self.non_printing_chars_re.sub("", text) |
|
|
