Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,67 +1,15 @@
-import streamlit as st
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# ==============================
-# PAGE CONFIG
-# ==============================
-st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
-
-# ==============================
-# LOAD MODEL
-# ==============================
-@st.cache_resource
-def load_model():
-    model_name = "deepseek-ai/deepseek-coder-6.7b-instruct"   # 🔥 BEST MODEL
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32,
-        device_map="auto"   # auto uses GPU if available
-    )
-
-    return tokenizer, model
-
-tokenizer, model = load_model()
-
-# ==============================
-# CLEAN OUTPUT
-# ==============================
-def clean_code(code):
-    code = code.strip()
-
-    remove_words = [
-        "Explanation:",
-        "Here is the code:",
-        "Output:",
-        "Answer:"
-    ]
-
-    for word in remove_words:
-        code = code.replace(word, "")
-
-    # Remove markdown formatting
-    code = code.replace("```python", "").replace("```", "")
-
-    return code.strip()
-
-# ==============================
-# GENERATE CODE
-# ==============================
 def generate_code(prompt, language):
 
     full_prompt = f"""
-You are an expert {language} developer.
+You are an expert {language} programmer.
 
-Generate clean, correct, and complete code.
+Write clean, correct, and complete code.
 
-Rules:
-- Only return code
-- No explanation
+Requirements:
+- No explanations
+- Only code
 - Proper syntax
-- Complete working solution
+- Complete function/program
 
 Task:
 {prompt}
@@ -74,59 +22,19 @@ Code:
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=350,
-            do_sample=True,
-            temperature=0.2,
-            top_p=0.9,
-            repetition_penalty=1.1,
+            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)
-
-# ==============================
-# SESSION STATE (CHAT HISTORY)
-# ==============================
-if "history" not in st.session_state:
-    st.session_state.history = []
-
-# ==============================
-# UI
-# ==============================
-st.title("💻 AI Code Generator (Advanced)")
-
-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 high-quality code..."):
-            code = generate_code(user_prompt, language)
-
-        st.session_state.history.append((user_prompt, code))
-
-# ==============================
-# DISPLAY HISTORY
-# ==============================
-for q, c in reversed(st.session_state.history):
-    st.markdown(f"### 🧑 Task:\n{q}")
-    st.code(c, language=language.lower())
+    return clean_code(result)