Update superposed/llama/tokenizer.py

superposed/llama/tokenizer.py

@@ -35,7 +35,7 @@ class Tokenizer:
         )
         assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
 
-    def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
+    def encode(self, s, bos: bool, eos: bool):
         """
         Encodes a string into a list of token IDs.
 
@@ -47,15 +47,16 @@ class Tokenizer:
         Returns:
             List[int]: A list of token IDs.
         """
-        assert type(s) is str
         t = self.sp_model.encode(s)
         if bos:
-            t = [self.bos_id] + t
+            for i in range(len(t)):
+                t[i] = [self.bos_id] + t[i]
         if eos:
-            t = t + [self.eos_id]
+            for i in range(len(t)):
+                t[i] = t[i] + [self.eos_id]
         return t
 
-    def decode(self, t: List[int]) -> str:
+    def decode(self, t):
         """
         Decodes a list of token IDs into a string.