Update app.py

app.py

@@ -1,28 +1,30 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+from transformers import MBartForConditionalGeneration, MBartTokenizer
 
-# Utilisation du modèle sdadas/byt5-text-correction
-model_name = "sdadas/byt5-text-correction"
+# Remplace par ton repo exact si besoin :
+model_name = "alice/mini/mBART_french_correction"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Chargement du tokenizer et du modèle
+tokenizer = MBartTokenizer.from_pretrained(model_name)
+model = MBartForConditionalGeneration.from_pretrained(model_name)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
 def correction_grammaticale(texte):
-    # Pour le français, on ajoute le préfixe "<fr>" devant le texte
-    input_text = "<fr> " + texte
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    # Tokenisation
+    inputs = tokenizer(texte, return_tensors="pt", max_length=512, truncation=True).to(device)
     
+    # Génération
     outputs = model.generate(
-        input_ids,
+        **inputs,
         max_length=512,
         num_beams=4,
         early_stopping=True
     )
     
+    # Décodage
     correction = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return correction
 
@@ -30,7 +32,7 @@ demo = gr.Interface(
     fn=correction_grammaticale,
     inputs=gr.Textbox(label="Texte à corriger"),
     outputs=gr.Textbox(label="Texte corrigé"),
-    title="Correcteur de Texte Français"
+    title="Correcteur MBART Français"
 )
 
 if __name__ == "__main__":