Fix chat script

Author: andrewdalpino

chat.py

@@ -10,7 +10,8 @@
 from colored import fore_rgb, style
 
 from src.nope_gpt.model import NoPEGPT
-from nope_gpt.memory import BufferWindowMemory
+from src.nope_gpt.tokenization import BaseTokenizer, ChatTokenizer
+from src.nope_gpt.memory import BufferWindowMemory
 
 DEFAULT_SYSTEM_MESSAGE = (
     "You're a helpful AI assistant named NoPE GPT. "
@@ -52,7 +53,9 @@ def main():
         args.checkpoint_path, map_location=args.device, weights_only=False
     )
 
-    tokenizer = checkpoint["tokenizer"]
+    tokenizer = BaseTokenizer.from_tiktoken(**checkpoint["tokenizer_args"])
+
+    tokenizer = ChatTokenizer(tokenizer)
 
     model = NoPEGPT(**checkpoint["model_args"])
 

src/nope_gpt/tokenization.py

@@ -146,6 +146,12 @@ class ChatTokenizer(ModelHubMixin):
 
     RESPONSE_HEADER = "<|im_start|>assistant\n"
 
+    @classmethod
+    def from_tiktoken(cls, name: str) -> Self:
+        """Instantiate a tokenizer from a pretrained tiktoken tokenizer."""
+
+        return cls(get_encoding(name))
+
     @classmethod
     def _from_pretrained(
         cls,