import os
import gradio as gr
from transformers import pipeline

# Récupérer le token d'authentification depuis la variable d'environnement
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
# Définir le modèle à utiliser
model_name = "salmapm/llama2_salma"

# Sélectionner le pipeline avec authentification
generator = pipeline("text-generation", model=model_name, use_auth_token=HUGGINGFACE_TOKEN)

# Fonction de génération de texte
def generate_text(prompt):
    results = generator(prompt, max_length=50)
    return results[0]['generated_text']

# Créer l'interface utilisateur avec Gradio
iface = gr.Interface(
    fn=generate_text,
    inputs="text",
    outputs="text",
    title="Générateur de Texte",
    description="Entrez un texte et laissez le modèle générer la suite."
)

# Lancer l'application
iface.launch()