#BS_app.py_04
#Training OK BUT Need new folder to save training session
#testing bloom1b training


import gradio as gr
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset, Dataset
from huggingface_hub import HfApi, HfFolder
import requests
from io import BytesIO

# Récupérer token depuis les variables d'environnement
hf_token = os.getenv("MisterAI_bigscience_bloom_560m")

# Configurer le token pour l'utilisation avec Hugging Face
if hf_token:
    HfFolder.save_token(hf_token)
else:
    raise ValueError("Le token Hugging Face n'est pas configuré. Assurez-vous qu'il est défini dans les variables d'environnement.")

# Chargement du modèle et du tokenizer
model_name = "MisterAI/bigscience_bloom-560m"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Fonction pour générer une réponse
def generate_response(input_text):
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=100)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Fonction pour créer un nouveau répertoire d'entraînement et sauvegarder le modèle
def create_training_folder_and_save_model(dataset_path, dataset_file, epochs, batch_size):
    # Récupération du chemin du répertoire de base
#    base_path = os.path.join(model_name, "FT")
    base_path = os.path.join(model_name, "resolve", "main", "FT")

    
    # Récupération des dossiers existants et incrémentation du numéro
#    existing_folders = [folder for folder in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, folder))]
    existing_folders = [folder for folder in os.listdir(base_path) if os.path.isdir(os.path.join(base_path, folder))]  
    new_folder_num = len(existing_folders) + 1
    new_folder_name = f"TRAIN{new_folder_num:02d}"

    # Création du nouveau répertoire d'entraînement
    new_folder_path = os.path.join(base_path, new_folder_name)
    os.makedirs(new_folder_path, exist_ok=True)

    # Affichage des logs
    print("Création du nouveau répertoire d'entraînement en cours...")
    print(f"Nouveau répertoire créé : {new_folder_path}")

#    # Chargement du dataset
#    if dataset_path.startswith("https://huggingface.co/datasets/"):
#        dataset = load_dataset('json', data_files={dataset_file: dataset_path})
#    else:
#        dataset = load_dataset('json', data_files={dataset_file: dataset_path})

        
        

    # Récupération du fichier à partir de l'URL fournie
    response = requests.get(f"{dataset_path}/resolve/main/{dataset_file}")
    dataset_lines = response.text.strip().split('\n')

    # Convertir les lignes en dictionnaires pour JSONL
    dataset_dict = [json.loads(line) for line in dataset_lines if line]

    # Créer un Dataset Hugging Face
    dataset = Dataset.from_dict({
        'question': [item['question'] for item in dataset_dict],
        'chosen': [item['chosen'] for item in dataset_dict]
    })

        


    # Préparation des données
    dataset = Dataset.from_dict(dataset[dataset_file])
    dataset = dataset.map(lambda x: tokenizer(x['question'] + ' ' + x['chosen'], truncation=True, padding='max_length'), batched=True)
    dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])

    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Configuration de l'entraînement
    training_args = TrainingArguments(
        output_dir=new_folder_path,  # Sauvegarder dans le nouveau répertoire
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        save_steps=10_000,
        save_total_limit=2,
        push_to_hub=False,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset
    )

    # Lancement de l'entraînement
    trainer.train()

    # Sauvegarde du modèle dans le nouveau répertoire
    trainer.save_model(new_folder_path)
    tokenizer.save_pretrained(new_folder_path)

    # Push vers Hugging Face Hub avec le bon repo_id
    api = HfApi()
    api.upload_folder(
        folder_path=new_folder_path,
        repo_id=model_name,
        repo_type="model",
        token=hf_token
    )

    return f"Fine-tuning terminé et modèle sauvegardé dans {new_folder_path}."

# Interface Gradio
with gr.Blocks() as demo:
    with gr.Tab("Chatbot"):
        chat_interface = gr.Interface(
            fn=generate_response,
            inputs="text",
            outputs="text",
            title="Chat avec le modèle",
            description="Tester Le Modèle"
        )

    with gr.Tab("Fine-Tuning"):
        with gr.Row():
            dataset_path = gr.Textbox(label="Chemin du dataset")
            dataset_file = gr.Textbox(label="Nom du fichier du dataset")
            epochs = gr.Number(label="Nombre d'époques", value=1)
            batch_size = gr.Number(label="Taille du batch", value=2)
        
        with gr.Row():
            fine_tune_button = gr.Button("Lancer le Fine-Tuning")

        with gr.Row():
            fine_tune_output = gr.Textbox(label="État du Fine-Tuning")

        fine_tune_button.click(
            create_training_folder_and_save_model,
            inputs=[dataset_path, dataset_file, epochs, batch_size],
            outputs=fine_tune_output
        )

# Lancement de la démo
if __name__ == "__main__":
    demo.launch()

# Je n'ai pas trouvé de documentation officielle sur la création de répertoires dans l'espace de travail Hugging Face.