import streamlit as st
from transformers import pipeline

# Load the model pipeline
@st.cache_resource()
def load_model():
    return pipeline("text2text-generation", model="Salesforce/codet5p-220m")

model = load_model()

# Streamlit UI
st.title("CodeCorrect AI")
st.subheader("AI-powered Code Autocorrect Tool")

code_input = st.text_area("Enter your code here:", height=200)

if st.button("Correct Code"):
    if code_input.strip():
        response = model(code_input, max_length=512)
        corrected_code = response[0]['generated_text']
        st.text_area("Corrected Code:", corrected_code, height=200)
    else:
        st.warning("Please enter some code.")