Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,3 +1,62 @@
+import streamlit as st
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# ==============================
+# PAGE CONFIG
+# ==============================
+st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
+
+st.title("💻 AI Code Generator (Stable Version)")
+
+# ==============================
+# LOAD MODEL (SAFE)
+# ==============================
+@st.cache_resource
+def load_model():
+    model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"  # ✅ HF Free Safe
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32
+    )
+
+    model.eval()
+    return tokenizer, model
+
+# Load with spinner (IMPORTANT)
+with st.spinner("🔄 Loading AI model... Please wait"):
+    tokenizer, model = load_model()
+
+st.success("✅ Model Loaded Successfully")
+
+# ==============================
+# CLEAN OUTPUT
+# ==============================
+def clean_code(code):
+    code = code.strip()
+
+    # Remove unwanted text
+    remove_words = [
+        "Explanation:",
+        "Here is the code:",
+        "Output:",
+        "Answer:"
+    ]
+
+    for word in remove_words:
+        code = code.replace(word, "")
+
+    # Remove markdown
+    code = code.replace("```python", "").replace("```", "")
+
+    return code.strip()
+
+# ==============================
+# GENERATE CODE
+# ==============================
 def generate_code(prompt, language):
 
     full_prompt = f"""
@@ -5,11 +64,10 @@ You are an expert {language} programmer.
 
 Write clean, correct, and complete code.
 
-Requirements:
+Rules:
+- Only return code
 - No explanations
-- Only code
-- Proper syntax
-- Complete function/program
+- Complete solution
 
 Task:
 {prompt}
@@ -19,22 +77,66 @@ Code:
 
     inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=300,
-            do_sample=True,              # ✅ IMPORTANT
-            temperature=0.2,             # ✅ slight creativity
-            top_p=0.9,                   # ✅ better coherence
-            repetition_penalty=1.1,      # ✅ avoid loops
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Remove prompt part safely
-    if "Code:" in result:
-        result = result.split("Code:")[-1]
-
-    return clean_code(result)
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=200,
+                do_sample=True,
+                temperature=0.3,
+                top_p=0.9,
+                repetition_penalty=1.1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        if "Code:" in result:
+            result = result.split("Code:")[-1]
+
+        return clean_code(result)
+
+    except Exception as e:
+        return f"# ERROR: {str(e)}"
+
+# ==============================
+# SESSION STATE
+# ==============================
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+# ==============================
+# UI INPUT
+# ==============================
+col1, col2 = st.columns(2)
+
+with col1:
+    user_prompt = st.text_area("Describe your task", height=200)
+
+with col2:
+    language = st.selectbox(
+        "Select Programming Language",
+        ["Python", "JavaScript", "SQL", "Java", "C++", "HTML", "CSS"]
+    )
+
+# ==============================
+# GENERATE BUTTON
+# ==============================
+if st.button("Generate Code"):
+    if not user_prompt.strip():
+        st.warning("⚠️ Please enter a task")
+    else:
+        with st.spinner("⚡ Generating code..."):
+            code = generate_code(user_prompt, language)
+
+        st.session_state.history.append((user_prompt, code))
+
+# ==============================
+# OUTPUT DISPLAY
+# ==============================
+if st.session_state.history:
+    st.subheader("📌 Generated Results")
+
+    for q, c in reversed(st.session_state.history):
+        st.markdown(f"**🧑 Task:** {q}")
+        st.code(c, language=language.lower())