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

# Charger le modèle et le tokenizer
language = "french"
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
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

# Exemple d'utilisation (test rapide)
example_text = "L'amour, cet élan mystérieux qui unit les âmes, est une lumière douce dans les ténèbres du quotidien..."

# Interface Gradio
interface = gr.Interface(
    fn=generate_summary,
    inputs=gr.Textbox(
        lines=10, 
        placeholder="Entrez un texte à résumer ou transformer en poème..."
    ),
    outputs=gr.Textbox(
        label="Texte généré"
    ),
    examples=[[example_text]],
    title="Génération de texte poétique ou résumé",
    description=(
        "Cette application utilise un modèle pré-entraîné BART pour transformer "
        "un texte d'entrée en une sortie poétique ou résumé. Entrez une phrase ou un paragraphe pour tester !"
    ),
)

# Lancer l'application
interface.launch()