Update app.py

app.py

@@ -5,7 +5,7 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, LlamaForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, Llama3ForCausalLM
 from peft import AutoPeftModelForCausalLM
 
 DESCRIPTION = """\
@@ -24,7 +24,7 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 model_name = "estrogen/c4ai-command-r7b-12-2024"
 
-model = LlamaForCausalLM.from_pretrained(
+model = Llama3ForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float16,
     trust_remote_code=True