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

# Charger le modèle et le tokenizer
model_name = "Moustapha91/bart_large_poetique-v02"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Fonction de génération ajustée
# Fonction de génération ajustée
def generate_summary(input_text):
    inputs = tokenizer(
        input_text,
        padding="max_length",
        max_length=512,  # Ajustez si nécessaire
        return_tensors="pt",
        truncation=True,
    )
    input_ids = inputs.input_ids.to(model.device)
    attention_mask = inputs.attention_mask.to(model.device)

    # Génération avec des paramètres personnalisés
    outputs = model.generate(
        input_ids=input_ids,
        attention_mask=attention_mask,
        max_length=512,
        num_beams=4,
        temperature=1.0,
        early_stopping=True,
        repetition_penalty=6.0,
        no_repeat_ngram_size=5,
    )
    output_str = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output_str


# Liste d'exemples enrichis
examples = [
    "L'amour, cet élan mystérieux qui unit les âmes, est une lumière douce dans les ténèbres du quotidien.",
    "Les tirailleurs sénégalais, ces soldats venus d'Afrique, ont marqué l'histoire par leur courage et leur sacrifice. Recrutés à partir de 1857 sous l'autorité coloniale française, ils provenaient de nombreux pays d'Afrique de l'Ouest.",
    "Il était une fois un garçon plein de rêves et d'espoir. Porté par une vision audacieuse, il rêvait d'unifier toute l'Afrique sous un même idéal de paix et de solidarité.",
    "L'eau, source de vie, est l'élément essentiel qui relie tous les êtres vivants. Elle coule dans les rivières, traverse les montagnes et nourrit la terre.",
    "Dans un petit village au cœur de l'Afrique, une vieille femme racontait les légendes des ancêtres à la lumière des étoiles, transmettant ainsi leur sagesse aux générations futures."
]

# Personnalisation de la description et de l'interface
title = "📝 Générateur de Texte Poétique et Résumé"
description = """
Bienvenue dans notre application interactive de génération de texte poétique ou résumé ! 🖋️  
Utilisez ce modèle pré-entraîné pour transformer des idées en textes élégants et expressifs.  
Ajoutez un paragraphe dans la zone ci-dessous et voyez la magie opérer !
"""
about = """
### À propos
Ce projet utilise le modèle **BART** de Facebook, connu pour ses capacités avancées en résumé et génération de texte.  
Créé pour inspirer, éduquer et honorer l'héritage de la littérature et des récits africains.  
Vous pouvez tester avec les exemples ci-dessus ou entrer vos propres idées.

---

📚 **Exemples suggérés :**
- Un texte sur l'amour.  
- Une introduction sur les tirailleurs sénégalais.  
- Un récit d'unité et d'espoir en Afrique.
"""

# Interface Gradio améliorée
interface = gr.Interface(
    fn=generate_summary,
    inputs=gr.Textbox(
        lines=10,
        placeholder="Entrez un texte à transformer en poème ou résumé ici...",
        label="🖊️ Texte à transformer"
    ),
    outputs=gr.Textbox(
        lines=10,
        label="✨ Texte généré"
    ),
    examples=examples,
    title=title,
    description=description,
    theme="compact",  # Thème compact et moderne
    article=about,  # Section 'À propos'
    allow_flagging="never",  # Désactive le flagging si inutile
)

# Lancer l'application
interface.launch()