Upload tokenizer.py

tokenizer.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Sequence, List
+
+import sentencepiece as spm
+
+
+def train_sentencepiece(
+    data_files: Sequence[str],
+    model_prefix: str = 'tokenizer',
+    vocab_size: int = 32000,
+    model_type: str = 'bpe',
+    character_coverage: float = 1.0,
+    byte_fallback: bool = True,
+    pad_id: int = 1,
+    unk_id: int = 0,
+    bos_id: int = 2,
+    eos_id: int = 3,
+) -> str:
+    data_files = [str(Path(p)) for p in data_files]
+    if not data_files:
+        raise ValueError('data_files is empty')
+    args = [
+        f"--input={','.join(data_files)}",
+        f'--model_prefix={model_prefix}',
+        f'--vocab_size={int(vocab_size)}',
+        f'--model_type={model_type}',
+        f'--character_coverage={character_coverage}',
+        f'--pad_id={pad_id}',
+        f'--unk_id={unk_id}',
+        f'--bos_id={bos_id}',
+        f'--eos_id={eos_id}',
+        '--hard_vocab_limit=false',
+        '--normalization_rule_name=nmt_nfkc',
+    ]
+    if byte_fallback:
+        args.append('--byte_fallback=true')
+    spm.SentencePieceTrainer.train(' '.join(args))
+    return f'{model_prefix}.model'
+
+
+class TokenizerWrapper:
+    def __init__(self, model_path: str):
+        model_path = str(Path(model_path))
+        if not Path(model_path).exists():
+            raise FileNotFoundError(model_path)
+        self.sp = spm.SentencePieceProcessor(model_file=model_path)
+        self.vocab_size = int(self.sp.vocab_size())
+        self.pad_id = self.sp.pad_id()
+        self.unk_id = self.sp.unk_id()
+        self.bos_id = self.sp.bos_id()
+        self.eos_id = self.sp.eos_id()
+        for name, val in [('pad', self.pad_id), ('unk', self.unk_id), ('bos', self.bos_id), ('eos', self.eos_id)]:
+            if val < 0:
+                raise ValueError(f'SentencePiece model missing <{name}>')
+
+    def encode(self, text: str, add_bos: bool = True, add_eos: bool = False) -> List[int]:
+        ids = list(self.sp.encode(text, out_type=int))
+        if add_bos:
+            ids = [self.bos_id] + ids
+        if add_eos:
+            ids = ids + [self.eos_id]
+        return ids
+
+    def encode_batch(self, texts: Sequence[str], add_bos: bool = True, add_eos: bool = False) -> List[List[int]]:
+        return [self.encode(t, add_bos=add_bos, add_eos=add_eos) for t in texts]
+
+    def decode(self, ids: Sequence[int]) -> str:
+        return self.sp.decode([int(i) for i in ids if int(i) != self.pad_id])
+
+    def save_config(self, path: str) -> None:
+        Path(path).write_text(json.dumps({
+            'vocab_size': self.vocab_size,
+            'pad_id': self.pad_id,
+            'unk_id': self.unk_id,
+            'bos_id': self.bos_id,
+            'eos_id': self.eos_id,
+        }, indent=2), encoding='utf-8')