Fixed an error regarding the verification of a type

chatbot.py

@@ -17,7 +17,7 @@ prompt.
 import argparse
 import torch
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase
 from typing import Optional, Union
 
 
@@ -28,9 +28,9 @@ class ChatBot:
 
     Parameters
     ----------
-    base_model : str | AutoModelForCausalLM
+    base_model : str | PreTrainedModel
         A name (path in hugging face hub) for a model, or the model itself.
-    tokenizer : AutoTokenizer | None
+    tokenizer : PreTrainedTokenizerBase | None
         Needed in case the base_model is a given model, otherwise it will load
         the same model given by the base_model path.
     initial_prompt : str
@@ -53,8 +53,8 @@ class ChatBot:
     """
 
     def __init__(self,
-                 base_model: Union[str, AutoModelForCausalLM],
-                 tokenizer: Optional[AutoTokenizer] = None,
+                 base_model: Union[str, PreTrainedModel],
+                 tokenizer: Optional[PreTrainedTokenizerBase] = None,
                  initial_prompt: Optional[str] = None,
                  keep_context: bool = False,
                  creative: bool = False,
@@ -69,7 +69,7 @@ class ChatBot:
             )
             self.tokenizer = AutoTokenizer.from_pretrained(base_model)
         else:
-            assert isinstance(self.tokenizer, AutoTokenizer),\
+            assert isinstance(tokenizer, PreTrainedTokenizerBase),\
                 "If the base model is given, the tokenizer should be given as well"
             self.model = base_model
             self.tokenizer = tokenizer