Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
 
 # ==============================
 # 🔐 HIDE STREAMLIT MENU
@@ -20,18 +21,31 @@ footer {visibility: hidden;}
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# LOAD MODEL (OPTIMIZED)
+# HF TOKEN (SECURE)
+# ==============================
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+if not HF_TOKEN:
+    st.error("❌ HF_TOKEN not found. Add it in Hugging Face Secrets.")
+    st.stop()
+
+# ==============================
+# LOAD MODEL (CACHED)
 # ==============================
 @st.cache_resource
 def load_model():
     model_name = "google/codegemma-2b"
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        token=HF_TOKEN
+    )
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float32,   # CPU safe
-        device_map="cpu"             # force CPU (faster on Spaces)
+        token=HF_TOKEN,
+        torch_dtype=torch.float16,
+        device_map="auto"
     )
 
     return tokenizer, model
@@ -56,25 +70,25 @@ Rules:
 - No explanation
 """
 
-    inputs = tokenizer(full_prompt, return_tensors="pt")
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
 
-    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
-        )
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=300,
+        temperature=0.2,
+        top_p=0.9,
+        do_sample=True
+    )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # Remove prompt from output
     return result.replace(full_prompt, "").strip()
 
 # ==============================
 # UI
 # ==============================
-st.title("💻 AI Code Generator")
+st.title("💻 AI Code Generator (CodeGemma)")
 
 col1, col2 = st.columns(2)
 
@@ -87,12 +101,19 @@ with col2:
         ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
     )
 
-if st.button("Generate Code"):
+# ==============================
+# GENERATE BUTTON
+# ==============================
+if st.button("🚀 Generate Code"):
     if not user_prompt.strip():
         st.warning("Please enter a task")
     else:
-        with st.spinner("⚡ Generating fast code..."):
-            code = generate_code(user_prompt, language)
+        with st.spinner("Generating code..."):
+            try:
+                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())
+            except Exception as e:
+                st.error(f"❌ Error: {str(e)}")