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