import gradio as gr
from gector.gec_model import GecBERTModel

# Load the GECToR model
def load_model():
    model = GecBERTModel(
        vocab_path='data/output_vocabulary',
        model_paths=['data/model_files/xlnet_0_gector.th'],
        max_len=128, min_len=3
    )
    return model

# Function to correct grammar using GECToR model
def correct_grammar(text):
    # Initialize the model (you can load it once and use globally to avoid reloading each time)
    model = load_model()
    # Correct the input text
    corrected_text = model.handle_batch([text])
    return corrected_text[0]  # Since the result is a list, return the first (and only) item

# Define Gradio interface
def create_gradio_interface():
    # Input and output in Gradio
    interface = gr.Interface(
        fn=correct_grammar,  # Function to run
        inputs="text",  # Input is plain text
        outputs="text",  # Output is plain text
        title="Grammar Correction App using GECToR",
        description="Enter your text, and this app will correct its grammar using GECToR."
    )
    return interface

# Launch the Gradio app
if __name__ == "__main__":
    gradio_interface = create_gradio_interface()
    gradio_interface.launch()