import gradio as gr
from transformers import pipeline
import torch

# Load the Grammarly CoEdit model for grammar & spelling correction
pipe = pipeline("text2text-generation", 
                model="grammarly/coedit-large",
                torch_dtype=torch.bfloat16)

def correct_grammar(text: str) -> str:
    """
    Takes an English sentence or paragraph and returns
    a grammatically and orthographically corrected version.
    """
    # Generate corrected text
    outputs = pipe(
        text,
        max_length=256,
        num_beams=5,
        early_stopping=True
    )
    return outputs[0]["generated_text"]

# Build a simple Gradio interface
with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown(
        """
        # 📝 Grammar & Spelling Corrector  
        Paste or type any English sentence below and receive a corrected version  
        powered by `grammarly/coedit-large`.
        """
    )

    with gr.Row():
        inp = gr.Textbox(
            label="✍️ Input Text",
            placeholder="Type a sentence or paragraph here…",
            lines=4,
            show_copy_button=True
        )
        out = gr.Textbox(
            label="✅ Corrected Text",
            placeholder="The corrected text will appear here.",
            lines=4,
            interactive=False,
            show_copy_button=True
        )

    btn = gr.Button("✔️ Correct Grammar")
    btn.click(fn=correct_grammar, inputs=inp, outputs=out)

    gr.Markdown(
        """
        ---  
        Built with 🤗 Transformers (`grammarly/coedit-large`) and 🚀 Gradio
        """
    )

demo.launch()