Spaces:
Sleeping
Sleeping
File size: 2,826 Bytes
c7f53d4 7087b82 c7f53d4 7087b82 ddf509e 7087b82 ddf509e 7087b82 45b1038 dd6b048 7087b82 ddf509e 7087b82 ddf509e 7087b82 ddf509e d144157 ddf509e d73111a d144157 ddf509e c7f53d4 45b1038 7087b82 ddf509e 7087b82 ddf509e d73111a ddf509e 7087b82 ddf509e 7087b82 5d9a347 ddf509e 5d9a347 ddf509e 7087b82 ddf509e 7087b82 ddf509e 5d9a347 ddf509e | 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 | import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# ==============================
# 🔐 HIDE STREAMLIT MENU
# ==============================
st.markdown("""
<style>
#MainMenu {visibility: hidden;}
header {visibility: hidden;}
footer {visibility: hidden;}
.stDeployButton {display:none;}
</style>
""", unsafe_allow_html=True)
# ==============================
# PAGE CONFIG
# ==============================
st.set_page_config(page_title="💻 AI Code Generator", layout="wide")
# ==============================
# LOAD MODEL (FAST + STABLE)
# ==============================
@st.cache_resource
def load_model():
model_name = "microsoft/phi-2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float32, # CPU friendly
device_map=None
)
return tokenizer, model
tokenizer, model = load_model()
# ==============================
# CODE GENERATION FUNCTION
# ==============================
def generate_code(prompt, language):
full_prompt = f"""### Instruction:
Write ONLY valid {language} code.
Do not include explanations.
Do not include special tokens.
Task:
{prompt}
### Response:
"""
inputs = tokenizer(full_prompt, return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=150,
temperature=0.2,
top_p=0.9,
do_sample=False
)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extract response
if "### Response:" in result:
result = result.split("### Response:")[-1]
# Remove unwanted tokens
unwanted_tokens = [
"<|endoftext|>",
"<|file_separator|>",
"<|assistant|>",
"<|system|>"
]
for token in unwanted_tokens:
result = result.replace(token, "")
return result.strip()
# ==============================
# UI
# ==============================
st.title("💻 AI Code Generator (Fast & Clean)")
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..."):
try:
code = generate_code(user_prompt, language)
st.success("✅ Generated Code")
st.code(code, language=language.lower())
except Exception as e:
st.error(f"❌ Error: {str(e)}") |