import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Modèle Hugging Face sélectionné
model_name = "PoloHuggingface/French_grammar_error_corrector"

# Chargement du modèle et du tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Vérification GPU/CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

def correction_grammaticale(texte):
    """
    Fonction qui envoie le texte au modèle T5 pour correction grammaticale.
    """
    # Préfixe facultatif (à tester si nécessaire)
    input_text = texte

    # Tokenisation
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
    
    # Génération du texte corrigé
    outputs = model.generate(
        **inputs,
        max_length=512,
        num_beams=4,
        early_stopping=True
    )
    
    # Décodage du texte corrigé
    correction = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return correction

# Interface utilisateur Gradio
demo = gr.Interface(
    fn=correction_grammaticale,
    inputs=gr.Textbox(label="Texte à corriger"),
    outputs=gr.Textbox(label="Texte corrigé"),
    title="Correcteur Grammatical Français"
)

if __name__ == "__main__":
    demo.launch()