Update Perceptrix/engine.py

Perceptrix/engine.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
+from transformers import LlamaForCausalLM, LlamaTokenizer, BitsAndBytesConfig, GenerationConfig
 from utils import setup_device
 import torch
 import tqdm
@@ -17,7 +17,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float32 if device == "cpu" else torch.bfloat16
 )
 
-model = AutoModelForCausalLM.from_pretrained(
+model = LlamaForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float32 if device == "cpu" else torch.bfloat16,
     device_map="auto",
@@ -27,7 +27,7 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config if str(device) != "cpu" else None,
 )
 
-tokenizer = AutoTokenizer.from_pretrained(
+tokenizer = LlamaTokenizer.from_pretrained(
     model_name,
     trust_remote_code=True,
     use_fast=True,