File size: 4,141 Bytes
65796a2
 
a4bcde0
65796a2
7067153
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65796a2
 
 
 
 
 
 
 
7067153
65796a2
 
 
7067153
 
 
 
 
 
65796a2
7067153
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a4bcde0
7067153
 
 
 
 
a4bcde0
 
 
 
 
 
7067153
 
a4bcde0
7067153
a4bcde0
 
 
7067153
 
 
 
 
a4bcde0
7067153
 
 
 
 
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
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
from datetime import datetime

# Custom CSS for a Grok/ChatGPT-like look
st.markdown("""
    <style>
    .main { background-color: #f9f9f9; padding: 20px; }
    .stTextArea textarea {
        border: 1px solid #ddd;
        border-radius: 8px;
        padding: 10px;
        font-family: 'Roboto', sans-serif;
        font-size: 16px;
        background-color: #fff;
        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
    }
    .stButton button {
        background-color: #4a90e2;
        color: white;
        border-radius: 8px;
        padding: 10px 20px;
        font-family: 'Roboto', sans-serif;
        font-size: 14px;
    }
    .stButton button:hover {
        background-color: #357abd;
    }
    .code-output {
        background-color: #2b2b2b;
        color: #f0f0f0;
        padding: 15px;
        border-radius: 8px;
        font-family: 'Courier New', monospace;
        font-size: 14px;
        margin-top: 10px;
    }
    .title {
        font-family: 'Roboto', sans-serif;
        font-size: 28px;
        font-weight: bold;
        color: #333;
        margin-bottom: 10px;
    }
    .subtitle {
        font-family: 'Roboto', sans-serif;
        font-size: 16px;
        color: #666;
        margin-bottom: 20px;
    }
    .chat-message {
        font-family: 'Roboto', sans-serif;
        font-size: 16px;
        color: #333;
        margin-bottom: 5px;
    }
    </style>
""", unsafe_allow_html=True)

# Load model and tokenizer
try:
    checkpoint = "Salesforce/codegen-350M-mono"
    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    model = AutoModelForCausalLM.from_pretrained(checkpoint)
except Exception as e:
    st.error(f"Error loading model: {e}")
    st.stop()

# Function to generate code
def generate_code(description):
    prompt = f"Generate Python code for the following task: {description}\n"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_length=500,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id
    )
    code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return code[len(prompt):].strip()

# Initialize chat history in session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# UI Layout
st.markdown('<div class="title">Code Generation Bot</div>', unsafe_allow_html=True)
st.markdown('<div class="subtitle">Describe your task, and I’ll generate Python code for you!</div>', unsafe_allow_html=True)

with st.container():
    # Input area
    description = st.text_area(
        "Enter your description here",
        placeholder="e.g., Write a function to calculate the factorial of a number",
        height=150
    )
    
    col1, col2 = st.columns([1, 1])  # Adjusted for two buttons
    with col1:
        if st.button("Generate"):
            if description.strip():
                with st.spinner("Thinking..."):
                    generated_code = generate_code(description)
                # Append to chat history with timestamp
                st.session_state.chat_history.append({
                    "input": description,
                    "output": generated_code,
                    "time": datetime.now().strftime("%H:%M:%S")
                })
            else:
                st.warning("Please enter a description first!")
    
    with col2:
        if st.button("Clear History"):
            st.session_state.chat_history = []
            st.success("Chat history cleared!")

    # Display chat history
    if st.session_state.chat_history:
        st.write("### Chat History")
        for chat in st.session_state.chat_history:
            st.markdown(f'<div class="chat-message"><strong>You ({chat["time"]}):</strong> {chat["input"]}</div>', unsafe_allow_html=True)
            st.markdown(f'<div class="code-output">{chat["output"]}</div>', unsafe_allow_html=True)
            st.markdown("---")  # Separator for readability

# Optional tip at the bottom
st.info("Tip: Check the generated code for accuracy before using it!")