Improve Tokenizer Class: Error Handling, Flexibility

llama/tokenizer.py

@@ -3,7 +3,7 @@
 
 import os
 from logging import getLogger
-from typing import List
+from typing import List, Union
 
 from sentencepiece import SentencePieceProcessor
 
@@ -12,25 +12,35 @@
 
 
 class Tokenizer:
-    def __init__(self, model_path: str):
-        # reload tokenizer
-        assert os.path.isfile(model_path), model_path
-        self.sp_model = SentencePieceProcessor(model_file=model_path)
-        logger.info(f"Reloaded SentencePiece model from {model_path}")
-
-        # BOS / EOS token IDs
-        self.n_words: int = self.sp_model.vocab_size()
-        self.bos_id: int = self.sp_model.bos_id()
-        self.eos_id: int = self.sp_model.eos_id()
-        self.pad_id: int = self.sp_model.pad_id()
-        logger.info(
-            f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
-        )
-        assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
-
-    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
-        assert type(s) is str
-        t = self.sp_model.encode(s)
+    def __init__(self, model_path: Union[str, None] = None):
+
+        if model_path is not None:
+            if not os.path.isfile(model_path):
+                raise FileNotFoundError(f"Model file not found: {model_path}")
+            self.sp_model = SentencePieceProcessor(model_file=model_path)
+            logger.info(f"Reloaded SentencePiece model from {model_path}")
+
+            # BOS / EOS / PAD / UNK token IDs
+            self.n_words: int = self.sp_model.vocab_size()
+            self.bos_id: int = self.sp_model.bos_id()
+            self.eos_id: int = self.sp_model.eos_id()
+            self.pad_id: int = self.sp_model.pad_id()
+            self.unk_id: int = self.sp_model.unk_id()
+            logger.info(
+                f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
+            )
+            assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
+
+    def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
+        assert isinstance(s, str), "Input 's' must be a string"
+        try:
+            t = self.sp_model.encode(s)
+        except Exception as e:
+            raise ValueError(f"Error during tokenization: {e}")
+
+        # Handle unknown tokens
+        t = [token_id if token_id in range(self.n_words) else self.unk_id for token_id in t]
+
         if bos:
             t = [self.bos_id] + t
         if eos: