trial / app.py
mishiawan's picture
Create app.py
94a22bf verified
import streamlit as st
from transformers import pipeline
# Title of the application
st.title("AI Code Bot")
# Introduction
st.write("""
### Generate Code Snippets with AI
Enter a programming-related prompt, and the AI will generate a code snippet for you.
""")
# Input box for user query
user_input = st.text_area("Enter your prompt (e.g., 'Write a Python function to reverse a list'):")
# Load the Hugging Face pipeline (code generation)
@st.cache_resource
def load_model():
# Authenticate automatically in Hugging Face Spaces
model = pipeline(
"text-generation",
model="EleutherAI/gpt-neo-1.3B" # The model to use
)
return model
model = load_model()
# When the user clicks the button
if st.button("Generate Code"):
if user_input.strip() == "":
st.warning("Please enter a prompt.")
else:
with st.spinner("Generating code..."):
try:
# Generate code using the model
result = model(user_input, max_length=100, num_return_sequences=1)
generated_code = result[0]["generated_text"]
# Display the output
st.code(generated_code, language="python")
except Exception as e:
st.error(f"An error occurred: {e}")