Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import os
 
 # ==============================
 # 🔐 HIDE STREAMLIT MENU
@@ -21,31 +20,18 @@ footer {visibility: hidden;}
 st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
 
 # ==============================
-# 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)
+# LOAD MODEL (FAST)
 # ==============================
 @st.cache_resource
 def load_model():
-    model_name = "google/codegemma-2b"
+    model_name = "microsoft/phi-2"
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        token=HF_TOKEN
-    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        token=HF_TOKEN,
-        torch_dtype=torch.float16,
-        device_map="auto"
+        torch_dtype=torch.float32,   # CPU friendly
+        device_map=None              # avoid GPU issues
     )
 
     return tokenizer, model
@@ -57,38 +43,37 @@ tokenizer, model = load_model()
 # ==============================
 def generate_code(prompt, language):
 
-    full_prompt = f"""
-You are an expert {language} developer.
-
-Write clean, optimized, production-ready code.
+    full_prompt = f"""### Instruction:
+Write a {language} program for the following task.
 
 Task:
 {prompt}
 
-Rules:
-- Only return code
-- No explanation
+### Response:
 """
 
-    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,
+        max_new_tokens=150,
         temperature=0.2,
         top_p=0.9,
-        do_sample=True
+        do_sample=False
     )
 
     result = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Remove prompt from output
-    return result.replace(full_prompt, "").strip()
+    # Extract only response part
+    if "### Response:" in result:
+        result = result.split("### Response:")[-1]
+
+    return result.strip()
 
 # ==============================
 # UI
 # ==============================
-st.title("💻 AI Code Generator (CodeGemma)")
+st.title("💻 AI Code Generator (Fast & Accurate)")
 
 col1, col2 = st.columns(2)