import streamlit as st
from transformers import pipeline

# Set up the Streamlit app layout and title
st.title("CodeWise - AI-based Code Commenting Tool")
st.subheader("Paste your code below, and the AI will generate comments explaining it!")

# Create an input text area where users can paste their code
code_input = st.text_area("Paste your code here:", height=300)

# Slider to select the length of comments
comment_length = st.slider("Select the comment length", min_value=50, max_value=300, value=150, step=10)

# Create a button to generate the comments
if st.button("Generate Comments"):
    if code_input.strip() == "":
        st.warning("Please paste some code before generating comments.")
    else:
        # Load a pre-trained model from Hugging Face for code summarization
        generator = pipeline("text2text-generation", model="Salesforce/codet5-base-multi-sum")

        # Generate comments using the model
        prompt = f"Comment this code:\n\n{code_input}"
        result = generator(prompt, max_length=comment_length, num_return_sequences=1)

        # Display the generated comments in the app
        st.subheader("Generated Comments:")
        st.code(result[0]['generated_text'], language='python')

# Display a footer message
st.write("Powered by [Hugging Face Transformers](https://huggingface.co/) and Streamlit.")