import gradio as gr
import os
import subprocess
from huggingface_hub import login

# Interface pour l'entraînement
def train_model(token, username, model_name, dataset_name, epochs, batch_size, learning_rate):
    # Authentification
    login(token=token)
    
    # Cloner le dépôt du modèle
    repo_id = f"{username}/{model_name}"
    os.system(f"git clone https://huggingface.co/{repo_id}")
    os.chdir(model_name)
    
    # Installer les dépendances
    os.system("pip install transformers datasets torch numpy tqdm wandb matplotlib pandas scikit-learn huggingface_hub accelerate evaluate tensorboard")
    
    # Créer un script d'entraînement
    train_script = f"""
from transformers import Trainer, TrainingArguments
from datasets import load_dataset
import torch
import torch.nn as nn
import os
import sys
import logging

# Configuration du logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Importer les modules nécessaires
from modeling_recipe_generator import RecipeGeneratorForHF, RecipeGeneratorConfig

def main():
    logger.info("Chargement du dataset depuis Hugging Face...")
    # Charger le dataset depuis Hugging Face
    dataset = load_dataset("{username}/{dataset_name}")
    
    logger.info("Chargement de la configuration du modèle...")
    # Charger la configuration
    config = RecipeGeneratorConfig.from_pretrained("{username}/{model_name}")
    
    logger.info("Création du modèle...")
    # Créer le modèle
    model = RecipeGeneratorForHF(config)
    
    logger.info("Configuration des arguments d'entraînement...")
    # Définir les arguments d'entraînement
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs={epochs},
        per_device_train_batch_size={batch_size},
        per_device_eval_batch_size={batch_size},
        warmup_steps=500,
        weight_decay=0.01,
        logging_dir="./logs",
        logging_steps=100,
        evaluation_strategy="steps",
        eval_steps=500,
        save_steps=1000,
        save_total_limit=2,
        learning_rate={learning_rate},
        fp16=True,  # Utiliser la précision mixte
        gradient_accumulation_steps=4,
        push_to_hub=True,
        hub_model_id="{username}/{model_name}",
        hub_strategy="every_save"
    )
    
    # Définir la fonction de calcul des métriques
    def compute_metrics(eval_pred):
        predictions, labels = eval_pred
        
        # Calculer la perplexité
        loss = torch.nn.functional.cross_entropy(
            torch.tensor(predictions).view(-1, predictions.shape[-1]),
            torch.tensor(labels).view(-1)
        )
        perplexity = torch.exp(loss)
        
        return {{
            "perplexity": perplexity.item()
        }}
    
    logger.info("Initialisation du Trainer...")
    # Initialiser le trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        compute_metrics=compute_metrics
    )
    
    logger.info("Démarrage de l'entraînement...")
    # Lancer l'entraînement
    trainer.train()
    
    logger.info("Entraînement terminé. Sauvegarde du modèle final...")
    # Sauvegarder le modèle final
    trainer.save_model("./final_model")
    
    logger.info("Envoi du modèle final vers Hugging Face Hub...")
    # Pousser le modèle final vers Hugging Face
    trainer.push_to_hub()
    
    logger.info("Processus d'entraînement terminé avec succès!")

if __name__ == "__main__":
    main()
"""
    
    with open("train_custom.py", "w") as f:
        f.write(train_script)
    
    # Lancer l'entraînement
    result = subprocess.run(["python", "train_custom.py"], capture_output=True, text=True)
    
    return result.stdout + "\n\n" + result.stderr

# Créer l'interface
iface = gr.Interface(
    fn=train_model,
    inputs=[
        gr.Textbox(label="Token Hugging Face", type="password"),
        gr.Textbox(label="Nom d'utilisateur Hugging Face"),
        gr.Textbox(label="Nom du modèle", value="recettes-generator"),
        gr.Textbox(label="Nom du dataset", value="recettes-dataset"),
        gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Nombre d'époques"),
        gr.Slider(minimum=1, maximum=32, value=8, step=1, label="Taille du batch"),
        gr.Slider(minimum=1e-6, maximum=1e-3, value=5e-5, step=1e-6, label="Taux d'apprentissage")
    ],
    outputs="text",
    title="Entraînement du modèle de génération de recettes",
    description="Entraînez votre modèle directement sur l'infrastructure de Hugging Face"
)

iface.launch()