Spaces:
Sleeping
Sleeping
File size: 3,491 Bytes
967d7c5 f521540 ddf509e f521540 fae2d0f d144157 fae2d0f 812346b 967d7c5 fae2d0f 967d7c5 812346b d73111a d144157 812346b c7f53d4 812346b 7087b82 967d7c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | 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"""
You are an expert {language} programmer.
Write clean, correct, and complete code.
Rules:
- Only return code
- No explanations
- Complete solution
Task:
{prompt}
Code:
"""
inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
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()) |