import gradio as gr
from transformers import pipeline

# Load the summarization model (facebook/bart-large-cnn)
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Use a smaller grammar correction model
grammar_correction_pipe = pipeline("text2text-generation", model="pszemraj/grammar-synthesis-small")

# Function for grammar correction
def correct_grammar(user_input):
    if user_input.strip():
        corrected_text = grammar_correction_pipe(user_input)[0]['generated_text']
        return corrected_text
    else:
        return "Please enter some text for grammar correction."

# Function for text summarization
def summarize_text(user_input):
    if user_input.strip():
        summary = summarizer(user_input, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
        return summary
    else:
        return "Please enter some text to summarize."

# Function to combine grammar correction and summarization
def correct_and_summarize(user_input):
    corrected_text = correct_grammar(user_input)  # First correct the grammar
    summary = summarize_text(corrected_text)  # Then summarize the corrected text
    return summary

# Gradio UI setup
with gr.Blocks() as demo:
    gr.Markdown("## Text Summarization and Grammar Correction Assistant")

    # Dropdown to select task
    task = gr.Dropdown(choices=["Summarize Text", "Correct Grammar"], label="Choose a task")

    # Input component for text
    user_input = gr.Textbox(label="Enter your text here:")

    # Output box for displaying the result
    output = gr.Textbox(label="Output", interactive=False)

    # Submit button
    submit_btn = gr.Button("Submit")

    # Function to process the input based on selected task
    def process_input(task, user_input):
        if task == "Summarize Text":
            return correct_and_summarize(user_input)  # Correct grammar, then summarize
        elif task == "Correct Grammar":
            return correct_grammar(user_input)  # Only correct grammar

    # Link the submit button to process the input
    submit_btn.click(process_input, inputs=[task, user_input], outputs=output)

# Launch the Gradio interface
demo.launch()