Upload model.py

model.py

@@ -0,0 +1,63 @@
+import os
+import kenlm
+import sentencepiece as spm
+from tokenizers import normalizers
+
+
+class KenlmModel:
+    def __init__(
+        self,
+        vocabulary_size: str,
+        ngram: str,
+        pruning: str,
+        normalize_nfd: bool = True,
+        normalize_numbers: bool = True,
+        normalize_puctuation: bool = True,
+    ):
+        self.model = kenlm.Model(os.path.join("files", f"jomleh-sp-{vocabulary_size}-o{ngram}-prune{pruning}.probing"))
+        self.tokenizer = spm.SentencePieceProcessor(os.path.join("files", f"jomleh-sp-{vocabulary_size}.model"))
+
+        norm_list = []
+        if normalize_numbers:
+            norm_list += [normalizers.Replace("۱", "۰"),
+                          normalizers.Replace("۲", "۰"),
+                          normalizers.Replace("۳", "۰"),
+                          normalizers.Replace("۴", "۰"),
+                          normalizers.Replace("۵", "۰"),
+                          normalizers.Replace("۶", "۰"),
+                          normalizers.Replace("۷", "۰"),
+                          normalizers.Replace("۸", "۰"),
+                          normalizers.Replace("۹", "۰"),
+                          normalizers.Replace(".", "")]
+        if normalize_puctuation:
+            norm_list += [normalizers.Replace(".", ""),
+                          normalizers.Replace("!", ""),
+                          normalizers.Replace("؛", ""),
+                          normalizers.Replace("،", ""),
+                          normalizers.Replace("؟", "")]
+        if normalize_nfd:
+            norm_list += [normalizers.NFD()]
+        norm_list += [normalizers.Strip()]
+
+        self.normalizer = normalizers.Sequence(norm_list)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        vocabulary_size: str,
+        ngram: str,
+        pruning: str,
+    ):
+        return cls(vocabulary_size, ngram, pruning)
+
+    def score(self, doc: str):
+        doc = self.normalizer.normalize_str(doc)
+        doc = ' '.join(self.tokenizer.encode(doc, out_type=str))
+        return self.model.score(doc)
+
+    def perplexity(self, doc: str):
+        doc = self.normalizer.normalize_str(doc)
+        doc = ' '.join(self.tokenizer.encode(doc, out_type=str))
+        log_score = self.model.score(doc)
+        length = len(doc.split()) + 1
+        return round(10.0 ** (-log_score / length), 1)