app.py

# Fichier : app.py
import os
import torch
import gradio as gr
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
from huggingface_hub import HfApi

# ==============================================================================
# 🛠️ Configuration et Modèles de Base Lamina
# ==============================================================================

# Dictionnaire des modèles de base disponibles (Modèle_ID: Nom_Affiché + Description)
MODELS_CONFIG = {
    "Finisha-LLM/Ilyana-pretrain": {
        "name": "Ilyana (fort en français)",
        "desc": "🚀 fort en français",
        "max_seq": 500 # Longueur de séquence ajustée pour ce petit modèle
    },
    "Finisha-LLM/Serena": {
        "name": "Serena, modèle espagnol",
        "desc": "✨ Bon compromis.",
        "max_seq": 328
    },
    "Finisha-LLM/Nelya-neko": {
        "name": "Nelya",
        "desc": "⚠️ Modèle nekolien.",
        "max_seq": 256
    }
}

# Clés pour le Dropdown Gradio
MODEL_CHOICES = list(MODELS_CONFIG.keys())


# ==============================================================================
# 🧠 Fonction Principale de Fine-Tuning
# ==============================================================================

def train_and_push_model(
    hf_token: str, 
    hf_username: str, 
    model_name: str, 
    num_epochs: int, 
    dataset_repo: str, 
    base_model_repo: str
):
    """
    Exécute le fine-tuning d'un modèle Lamina et le publie sur Hugging Face.
    """
    
    # --- 0. Initialisation et Connexion ---
    
    # Récupération de la configuration spécifique du modèle choisi
    config = MODELS_CONFIG[base_model_repo]
    max_seq_length = config["max_seq"]
    
    yield f"Début du processus pour le modèle de base : **{config['name']}**"
    
    # ⚠️ Vérification du Token et connexion
    if not hf_token or not hf_username:
        yield "❌ ERREUR : Le Token HF ou le Nom d'utilisateur est manquant. Veuillez les fournir."
        return
        
    os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
    api = HfApi()

    # Définition du nom final du repository sur le compte de l'utilisateur
    final_repo_name = f"{hf_username}/{model_name}"
    
    # --- 1. Chargement du Modèle et du Tokenizer ---
    yield f"Chargement du modèle : **{base_model_repo}**..."
    try:
        model = AutoModelForCausalLM.from_pretrained(base_model_repo)
        tokenizer = AutoTokenizer.from_pretrained(base_model_repo)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        yield "✅ Modèle et Tokenizer chargés."
    except Exception as e:
        yield f"❌ ERREUR lors du chargement du modèle/tokenizer : {e}"
        return

    # --- 2. Préparation du Dataset ---
    yield f"Chargement du dataset de Q&R : **{dataset_repo}**..."
    try:
        # Assurez-vous que le split 'train' est présent
        dataset = load_dataset(dataset_repo, split="train") 
        yield f"Dataset chargé. Nombre de lignes initiales : **{len(dataset)}**"
    except Exception as e:
        yield f"❌ ERREUR lors du chargement de la dataset '{dataset_repo}'. Vérifiez le nom/ID et sa disponibilité. Détail : {e}"
        return
    
    # Fonction de formatage et de tokenisation (celle que tu as fournie)
    def format_and_tokenize(example):
        # Utilisation de 'reponse' comme dans ton exemple
        formatted_text = f"### Instruction:\n{example['question']}\n\n### Response:\n{example['reponse']}{tokenizer.eos_token}"
        tokenized_output = tokenizer(
            formatted_text,
            truncation=True,
            padding="max_length",
            max_length=max_seq_length
        )
        tokenized_output["labels"] = tokenized_output["input_ids"].copy()
        return tokenized_output

    yield "Tokenisation et formatage du dataset en cours..."
    try:
        processed_dataset = dataset.map(format_and_tokenize, remove_columns=dataset.column_names)
        yield "✅ Dataset formaté et tokenisé."
    except Exception as e:
        yield f"❌ ERREUR lors du formatage de la dataset. Assurez-vous que les colonnes 'question' et 'reponse' existent. Détail : {e}"
        return


    # --- 3. Fine-Tuning ---
    output_dir = "./Lamina-FineTune-Temp"
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs, 
        per_device_train_batch_size=2, # Taille de batch par défaut raisonnable
        logging_steps=5,
        learning_rate=5e-5,
        save_total_limit=2,
        # Ajout du push vers le hub pour la fin
        push_to_hub=False # On va le faire manuellement après l'entraînement
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=processed_dataset,
        tokenizer=tokenizer,
    )

    yield f"🚀 Lancement du fine-tuning pour **{num_epochs}** époques..."
    try:
        trainer.train()
        yield "✅ Fine-tuning terminé !"
    except Exception as e:
        yield f"❌ ERREUR lors de l'entraînement. Cela peut être dû à un manque de ressources ou une erreur dans le dataset. Détail : {e}"
        return

    # --- 4. Sauvegarde et Publication sur Hugging Face ---
    
    # Sauvegarde locale du modèle entraîné
    final_model_path = "./Lamina-Trained-Final"
    trainer.save_model(final_model_path)
    tokenizer.save_pretrained(final_model_path)

    yield f"📦 Modèle sauvegardé localement. Publication sur **{final_repo_name}** en cours..."
    
    try:
        # Création ou utilisation d'un repository existant
        api.create_repo(repo_id=final_repo_name, exist_ok=True)
        # Téléchargement des fichiers
        api.upload_folder(
            folder_path=final_model_path,
            repo_id=final_repo_name,
            repo_type="model",
            commit_message=f"Fine-tune de {base_model_repo} par le Lamina Generator"
        )
        
        final_link = f"https://huggingface.co/{final_repo_name}"
        yield f"🎉 **SUCCÈS !** Votre modèle Lamina a été publié !\n"
        yield f"Consultez votre modèle ici : {final_link}\n"
        yield f"**Prochaine étape :** Créez un Space de démonstration pour votre nouveau modèle !"
        
    except Exception as e:
        yield f"❌ ERREUR lors de la publication sur Hugging Face. Vérifiez votre token (autorisation d'écriture) et votre nom d'utilisateur. Détail : {e}"
        return


# ==============================================================================
# 🖥️ Interface Gradio
# ==============================================================================

# Description détaillée pour l'interface
description = """
# 🧠 Lamina Generator : Créez votre propre SLM Lamina sans code !
Bienvenue Clemylia Community ! Cet outil vous permet de *fine-tuner* un de mes modèles Lamina de *pretrain* (lamina-basic, Small-lamina-pretrain, lamina-suite-pretrain) avec votre propre dataset de questions/réponses.
Votre nouveau modèle sera automatiquement publié sur votre compte Hugging Face !

---

**⚠️ Avertissements Importants :**

1.  **Token Nécessaire :** Vous devez utiliser un **Hugging Face Token avec les permissions d'écriture** (`write`) pour que le modèle puisse être publié sur votre compte.
2.  **Dataset :** Votre dataset doit être **publique** et contenir au moins deux colonnes nommées : `question` et `reponse`.
3.  **Taille :** Le temps d'entraînement dépend du modèle de base choisi et de la taille de votre dataset. Le grand modèle (`lamina-suite-pretrain`) peut prendre **plusieurs heures**.

Choisissez votre modèle de base ci-dessous pour voir la description complète et les recommandations.
"""

with gr.Blocks(title="Lamina Generator") as demo:
    gr.Markdown(description)

    # Ajout d'un affichage dynamique pour les descriptions de modèle
    model_info_output = gr.Markdown("Sélectionnez un modèle de base ci-dessous.")
    
    def update_model_info(model_id):
        if model_id in MODELS_CONFIG:
            config = MODELS_CONFIG[model_id]
            return f"**Modèle sélectionné :** {config['name']} \n\n**Recommandation :** {config['desc']}"
        return "Sélectionnez un modèle de base ci-dessous."
    
    with gr.Row():
        base_model_dropdown = gr.Dropdown(
            label="1. Choisir le Modèle de Base Lamina",
            choices=MODEL_CHOICES,
            value=MODEL_CHOICES[1], # Default to Small-lamina-pretrain
            interactive=True
        )
        base_model_dropdown.change(update_model_info, inputs=base_model_dropdown, outputs=model_info_output)
    
    # Affichage des infos/avertissements du modèle
    gr.Markdown("---")
    model_info_output = gr.Markdown(update_model_info(MODEL_CHOICES[1]))
    gr.Markdown("---")
    
    
    with gr.Accordion("🔑 Informations d'Authentification (Nécessaires pour publier)", open=True):
        hf_token_input = gr.Textbox(
            label="2. Votre Hugging Face Write Token", 
            type="password", 
            placeholder="hf_xxxxxxxxxxxxxxxxxxxxxx", 
            info="Token avec permission d'écriture pour la publication du modèle."
        )
        hf_username_input = gr.Textbox(
            label="3. Votre Nom d'utilisateur Hugging Face",
            placeholder="Clemylia",
            info="L'utilisateur sous lequel le modèle sera publié."
        )

    with gr.Accordion("📚 Configuration du Fine-Tuning", open=True):
        dataset_repo_input = gr.Textbox(
            label="4. Nom/ID Complet de votre Dataset (Hugging Face)",
            placeholder="votre_nom_utilisateur/nom_de_votre_dataset",
            info="Le dataset doit contenir les colonnes 'question' et 'reponse'."
        )
        model_name_input = gr.Textbox(
            label="5. Nom du Modèle Final",
            placeholder="mon-super-lamina",
            info="Le modèle sera publié sous le nom : <votre_nom>/<nom_du_modèle>."
        )
        epochs_input = gr.Slider(
            minimum=1, 
            maximum=100, 
            value=3, 
            step=1, 
            label="6. Nombre d'Époques d'Entraînement", 
            info="Plus d'époques = meilleure mémorisation, mais risque de 'sur-apprentissage'."
        )

    # Zone de Log et Bouton
    output_log = gr.Markdown("Prêt à entraîner ! Remplissez les champs et cliquez sur le bouton.")
    
    train_button = gr.Button("🤖 Entraîner mon propre Lamina ! 🚀", variant="primary")

    # Lien du bouton à la fonction Python
    train_button.click(
        fn=train_and_push_model,
        inputs=[
            hf_token_input, 
            hf_username_input, 
            model_name_input, 
            epochs_input, 
            dataset_repo_input, 
            base_model_dropdown
        ],
        outputs=output_log
    )
    
if __name__ == "__main__":
    demo.launch()