Update utils/inference.py

utils/inference.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import LlamaTokenizer, LlamaForCausalLM
 from peft import PeftModel
 from typing import Iterator
 from variables import SYSTEM, HUMAN, AI
@@ -8,7 +8,6 @@ from variables import SYSTEM, HUMAN, AI
 def load_tokenizer_and_model(base_model, adapter_model, load_8bit=True):
     """
     Loads the tokenizer and chatbot model.
-
     Args:
         base_model (str): The base model to use (path to the model).
         adapter_model (str): The LoRA model to use (path to LoRA model).
@@ -24,15 +23,15 @@ def load_tokenizer_and_model(base_model, adapter_model, load_8bit=True):
             device = "mps"
     except:  
         pass
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
+    tokenizer = LlamaTokenizer.from_pretrained(base_model)
     if device == "cuda":
-        model = AutoModelForCausalLM.from_pretrained(
+        model = LlamaForCausalLM.from_pretrained(
             base_model,
             load_in_8bit=load_8bit,
             torch_dtype=torch.float16
         )
     elif device == "mps":
-        model = AutoModelForCausalLM.from_pretrained(
+        model = LlamaForCausalLM.from_pretrained(
             base_model,
             device_map={"": device}
         )
@@ -44,7 +43,7 @@ def load_tokenizer_and_model(base_model, adapter_model, load_8bit=True):
                 torch_dtype=torch.float16,
             )
     else:
-        model = AutoModelForCausalLM.from_pretrained(
+        model = LlamaForCausalLM.from_pretrained(
             base_model, 
             device_map={"": device}, 
             low_cpu_mem_usage=True, 
@@ -76,7 +75,7 @@ shared_state = State()
 def decode(
     input_ids: torch.Tensor,
     model: PeftModel,
-    tokenizer: AutoTokenizer,
+    tokenizer: LlamaTokenizer,
     stop_words: list,
     max_length: int,
     temperature: float = 1.0,