Spaces:
Running
Running
File size: 4,542 Bytes
65796a2 a4bcde0 65796a2 c7a8ac7 7067153 c7a8ac7 65796a2 c7a8ac7 65796a2 7067153 65796a2 c7a8ac7 7067153 c7a8ac7 7067153 c7a8ac7 7067153 a4bcde0 7067153 a4bcde0 7067153 a4bcde0 7067153 a4bcde0 7067153 c7a8ac7 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 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
from datetime import datetime
# Custom CSS for UI
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)
# Cache model and tokenizer to avoid reloading
@st.cache_resource
def load_model_and_tokenizer():
checkpoint = "Salesforce/codegen-350M-mono"
try:
st.write("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
st.write("Loading model...")
model = AutoModelForCausalLM.from_pretrained(checkpoint)
st.write("Model and tokenizer loaded successfully!")
return tokenizer, model
except Exception as e:
st.error(f"Failed to load model/tokenizer: {e}")
return None, None
# Load model and tokenizer once
tokenizer, model = load_model_and_tokenizer()
if tokenizer is None or model is None:
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")
try:
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()
except Exception as e:
st.error(f"Error generating code: {e}")
return "Error: Could not generate code."
# Initialize chat history
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])
with col1:
if st.button("Generate"):
if description.strip():
with st.spinner("Thinking..."):
generated_code = generate_code(description)
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("---")
st.info("Tip: Check the generated code for accuracy before using it!") |