Update app.py

app.py

@@ -2,7 +2,9 @@ import os
 import gc
 import torch
 import gradio as gr
-from transformers import LlamaTokenizer, LlamaForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import StoppingCriteria, StoppingCriteriaList
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
 from huggingface_hub import login
 import os
 
@@ -24,7 +26,7 @@ print(f"Loading model from {MODEL_PATH} on {device}...")
 # =============================
 # Load Tokenizer and Model
 # =============================
-tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, legacy=False)
 model = LlamaForCausalLM.from_pretrained(
     MODEL_PATH,
     device_map="auto",