File size: 2,270 Bytes
c7f53d4
7087b82
 
c7f53d4
7087b82
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c7f53d4
7087b82
 
c7f53d4
7087b82
 
 
c7f53d4
 
7087b82
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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
# ==============================
@st.cache_resource
def load_model():
    model_name = "codellama/CodeLlama-7b-Instruct-hf"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto"
    )

    return tokenizer, model

tokenizer, model = load_model()

# ==============================
# CODE GENERATION FUNCTION
# ==============================
def generate_code(prompt, language):

    full_prompt = f"""
You are an expert {language} developer.

Write clean, optimized, production-ready code.

Task:
{prompt}

Rules:
- Only return code
- No explanation
"""

    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)

    outputs = model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.2,
        top_p=0.9
    )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return result.replace(full_prompt, "").strip()

# ==============================
# UI
# ==============================
st.title("πŸ’» AI Code Generator")

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"]
    )

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.success("βœ… Generated Code")
            st.code(code, language=language.lower())