Update app.py

app.py

@@ -1,9 +1,8 @@
-from transformers import LlamaForCausalLM, CodeLlamaTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load pre-trained model and tokenizer
-model_name = "codellama/CodeLlama-7b-hf"  # You can replace it with your model name
-tokenizer = CodeLlamaTokenizer.from_pretrained(model_name)
-model = LlamaForCausalLM.from_pretrained(model_name)
+# Load model directly
+tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-hf")
+model = AutoModelForCausalLM.from_pretrained("codellama/CodeLlama-7b-hf")
 
 # System message
 system_message = "You are a code teaching assistant named OmniCode created by Anusha K. Answer all the code related questions being asked."