import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the model and tokenizer from Hugging Face
model_name = "Salesforce/codet5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Streamlit UI
st.title("Code Generator")
st.write("Generate code snippets from natural language prompts using CodeT5!")

# Input for natural language prompt
prompt = st.text_area("Enter your coding task:", placeholder="Write a Python function to calculate the factorial of a number.")

# Slider to control output length
max_length = st.slider("Maximum length of generated code:", 20, 200, 50)

# Button to trigger code generation
if st.button("Generate Code"):
    if prompt.strip():
        # Tokenize and generate code
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
        outputs = model.generate(inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True)
        generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)

        # Display generated code
        st.write("### Generated Code:")
        st.code(generated_code, language="python")
    else:
        st.warning("Please enter a prompt to generate code.")