Update app.py

app.py

@@ -6,7 +6,7 @@ from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 # Load the model using Accelerate for memory optimization
 @st.cache_resource()
 def load_model():
-    MODEL_NAME = "codellama/CodeLlama-7b-hf"  # Model name
+    MODEL_NAME = "huggingface/CodeGen-2B"  # Updated model name
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     
     # Load model with accelerate to optimize for memory usage
@@ -30,8 +30,8 @@ def load_model():
 code_generator = load_model()
 
 # Streamlit UI
-st.title("CodeLlama-7B Code Bot 🚀")
-st.subheader("Generate code snippets using CodeLlama-7b-hf optimized for CPU")
+st.title("CodeGen-2B Code Bot 🚀")
+st.subheader("Generate code snippets using Hugging Face CodeGen-2B")
 
 # User input
 prompt = st.text_area("Enter a coding prompt (e.g., 'Write a Python function to sort a list'): ")
@@ -41,7 +41,7 @@ if st.button("Generate Code"):
     if prompt.strip():
         st.info("Generating code... Please wait ⏳")
         try:
-            # Generate code using the CodeLlama model
+            # Generate code using the CodeGen-2B model
             response = code_generator(
                 prompt,
                 max_length=512,  # Increase for longer code generation
@@ -57,4 +57,4 @@ if st.button("Generate Code"):
     else:
         st.warning("Please enter a prompt.")
 
-st.caption("Powered by CodeLlama-7B | Streamlit UI | CPU Optimized")
+st.caption("Created by Shamil")