Update utils/inference.py

utils/inference.py

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