Update app.py

app.py

@@ -1,37 +1,23 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-# Load the model using Accelerate for memory optimization
+# Load the model and tokenizer
 @st.cache_resource()
 def load_model():
-    MODEL_NAME = "Salesforce/codegen-2B-multi"  # Updated model name
+    MODEL_NAME = "Salesforce/codegen-350M-mono"  # Use a known model
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    
-    # 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
-        )
-    
-    # Use 'auto' or 'balanced' device_map for CPU or GPU
-    model = load_checkpoint_and_dispatch(
-        model,
-        MODEL_NAME,
-        device_map="auto",  # Automatically determine the best device allocation
-    )
-    
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float32)
+
+    # Load model into the pipeline for generation
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Initialize the model
 code_generator = load_model()
 
 # Streamlit UI
-st.title("CodeGen-2B Code Bot 🚀")
-st.subheader("Generate code snippets using Hugging Face CodeGen-2B")
+st.title("CodeGen Code Bot 🚀")
+st.subheader("Generate code snippets using Hugging Face CodeGen")
 
 # User input
 prompt = st.text_area("Enter a coding prompt (e.g., 'Write a Python function to sort a list'): ")
@@ -41,10 +27,10 @@ if st.button("Generate Code"):
     if prompt.strip():
         st.info("Generating code... Please wait ⏳")
         try:
-            # Generate code using the CodeGen-2B model
+            # Generate code using the model
             response = code_generator(
                 prompt,
-                max_length=512,  # Increase for longer code generation
+                max_length=512,  # Adjust for longer code if needed
                 temperature=0.2,  # Lower temperature for more deterministic results
                 do_sample=True,   # Enable sampling
                 num_return_sequences=1
@@ -57,4 +43,4 @@ if st.button("Generate Code"):
     else:
         st.warning("Please enter a prompt.")
 
-st.caption("Powered by CodeGen-2B | Streamlit UI | CPU Optimized")
+st.caption("Powered by CodeGen | Streamlit UI")