Update src/streamlit_app.py

src/streamlit_app.py

@@ -20,17 +20,18 @@ footer {visibility: hidden;}
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL
+# LOAD MODEL (OPTIMIZED)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "codellama/CodeLlama-7b-Instruct-hf"
+    model_name = "google/codegemma-2b"
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float16,
-        device_map="auto"
+        torch_dtype=torch.float32,   # CPU safe
+        device_map="cpu"             # force CPU (faster on Spaces)
     )
 
     return tokenizer, model
@@ -55,14 +56,16 @@ Rules:
 - No explanation
 """
 
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt")
 
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=300,
-        temperature=0.2,
-        top_p=0.9
-    )
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=150,   # reduced for speed
+            temperature=0.1,
+            top_p=0.85,
+            do_sample=True
+        )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
@@ -88,8 +91,8 @@ if st.button("Generate Code"):
     if not user_prompt.strip():
         st.warning("Please enter a task")
     else:
-        with st.spinner("Generating code..."):
+        with st.spinner("⚡ Generating fast code..."):
             code = generate_code(user_prompt, language)
 
-            st.success("✅ Generated Code")
-            st.code(code, language=language.lower())
+        st.success("✅ Generated Code")
+        st.code(code, language=language.lower())