Update app.py

app.py

@@ -1,25 +1,37 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-# Load the CodeLlama model and tokenizer
+# Load the model using Accelerate for memory optimization
 @st.cache_resource()
 def load_model():
     MODEL_NAME = "codellama/CodeLlama-7b-hf"  # Model name
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME, 
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        device_map="cpu"            # Map model to CPU
+    
+    # Load model with accelerate to optimize for memory usage
+    with init_empty_weights():
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME, 
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            low_cpu_mem_usage=True      # Enable low memory usage on CPU
+        )
+    
+    # Move model to CPU
+    model = load_checkpoint_and_dispatch(
+        model,
+        MODEL_NAME,
+        device_map="cpu",  # Load model onto CPU
     )
+    
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Initialize pipeline
+# Initialize the model
 code_generator = load_model()
 
 # Streamlit UI
 st.title("CodeLlama-7B Code Bot 🚀")
-st.subheader("Generate code snippets using CodeLlama-7b-hf on CPU")
+st.subheader("Generate code snippets using CodeLlama-7b-hf optimized for CPU")
 
 # User input
 prompt = st.text_area("Enter a coding prompt (e.g., 'Write a Python function to sort a list'): ")
@@ -38,7 +50,7 @@ if st.button("Generate Code"):
                 num_return_sequences=1
             )
             generated_code = response[0]['generated_text']
-            # Display the code output
+            # Display the generated code output
             st.code(generated_code, language="python")  # Change language as needed
         except Exception as e:
             st.error(f"Error: {str(e)}")