app.py

# BS_app.py_13
# Training OK - Correct Features BBAI
# training bigscience_bloom-560m
##HFChat Coder


# BS_app.py_11
# Training OK - Correct Features  BBAI
# training bigscience_bloom-560m

import gradio as gr
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from transformers.trainer_callback import TrainerCallback    
from datasets import load_dataset, Dataset
from huggingface_hub import HfApi, HfFolder
import requests
from io import BytesIO
import json
import psutil  # Ajout d'import psutil
import time  # Ajout d'import time pour ralentir l'entraînement
import threading  # Ajout d'import threading
import torch

# 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 surveiller l'utilisation des ressources GPU
def get_gpu_usage(gpu_limit):
    if gpu_limit > 0 and torch.cuda.is_available():
        torch.cuda.empty_cache()
        gpu_usage = (torch.cuda.memory_allocated() / torch.cuda.memory_total()) * 100
    else:
        gpu_usage = 0
    return gpu_usage

# Fonction pour surveiller l'utilisation des ressources CPU RAM
def monitor_resources(cpu_limit=90, ram_limit=90, gpu_limit=0):
    cpu_usage = psutil.cpu_percent()
    ram_usage = psutil.virtual_memory().percent
    gpu_usage = 0  # Placeholder pour GPU, à implémenter selon la bibliothèque utilisée
    return cpu_usage, ram_usage, gpu_usage

# Variable globale pour contrôler l'exécution du thread
thread = None


    
# Fonction pour le fine-tuning
def fine_tune_model(dataset_path, dataset_file, split, colonne_input_ids, colonne_attention_mask, colonne_labels, epochs, batch_size, num_cpus, ram_size, cpu_limit, ram_limit, gpu_limit):
    global thread
    if thread is None:
        thread = threading.Thread(target=fine_tune_model)
    thread.do_run = True
    logs = []

    # 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')
    dataset_dict = [json.loads(line) for line in dataset_lines if line]

    # Créer un Dataset Hugging Face
    dataset = Dataset.from_dict({
        colonne_input_ids: [item[colonne_input_ids] for item in dataset_dict],
        colonne_attention_mask: [item.get(colonne_attention_mask, "") for item in dataset_dict],
        colonne_labels: [item.get(colonne_labels, "") for item in dataset_dict]
    })

    # Préparation des données
    def preprocess_function(examples):    
        inputs = [inids + '' + atmask for inids, atmask in zip(examples[colonne_input_ids], examples[colonne_attention_mask])]    
        model_inputs = tokenizer(inputs, truncation=True, padding='max_length', max_length=2048)    
        model_inputs["labels"] = model_inputs["input_ids"].copy()    
        return model_inputs    
    
    dataset = dataset.map(preprocess_function, 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="./",
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        save_steps=10_000,
        save_total_limit=2,
        push_to_hub=True,  # Activer/Désactiver le Push > définir hub_model_id
        hub_model_id=model_name,  # Nom du dépôt existant ou nouveau dépot
        hub_token=hf_token,
    )

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

    # Lancement de l'entraînement avec logs
    for epoch in range(epochs):
        if not thread.do_run:
            logs.append("Arrêt demandé, sortie de l'entraînement.")
            break
        logs.append(f"Début de l'époque {epoch + 1}/{epochs}")

        # Surveiller les ressources pendant l'entraînement
        cpu_usage, ram_usage, gpu_usage = monitor_resources(cpu_limit, ram_limit, gpu_limit)
        logs.append(f"Utilisation CPU: {cpu_usage}%, Utilisation RAM: {ram_usage}%, Utilisation GPU: {gpu_usage}%")

        # Si les limites sont atteintes, ralentir l'entraînement
        if cpu_usage >= cpu_limit or ram_usage >= ram_limit or gpu_usage >= gpu_limit:
            logs.append(f"Limites atteintes : CPU {cpu_usage}%, RAM {ram_usage}%, GPU {gpu_usage}%. Ralentissement de l'entraînement.")
            time.sleep(10)  # Attendre 10 secondes avant de continuer

        trainer.train()
        logs.append(f"Fin de l'époque {epoch + 1}/{epochs}")

    return "Fine-tuning terminé et modèle sauvegardé.", logs

# Fonction pour arrêter l'entraînement
def stop_fine_tune():
    global thread
    if thread and thread.is_alive():
        thread.do_run = False
        thread.join()

# 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=f"Entrez votre message pour obtenir une réponse de {model_name}"  # Utilisation de model_name au lieu de modele_name
        )

    with gr.Tab("Fine-Tuning"):
        with gr.Row():
            num_cpus = gr.Number(label="Nombre de CPU", value=2)
            ram_size = gr.Number(label="Taille de RAM (en Go)", value=16)
            num_gpus = gr.Number(label="Nombre de GPU", value=0)
            cpu_limit = gr.Slider(minimum=0, maximum=100, label="Limite CPU (%)", value=75)
            ram_limit = gr.Slider(minimum=0, maximum=100, label="Limite RAM (%)", value=75)    
            gpu_limit = gr.Slider(minimum=0, maximum=100, label="Limite GPU (%)", value=0)    

        with gr.Row():
            dataset_path = gr.Textbox(label="Chemin du dataset", value="https://huggingface.co/datasets/MisterAI/SimpleSmallFrenchQA")
            dataset_file = gr.Textbox(label="Nom du fichier du dataset", value="Dataset02_01L_QR_256_Francais.jsonl")

        with gr.Row():
            split = gr.Textbox(label="Split (si applicable)", value="train")
            colonne_input_ids = gr.Textbox(label="Colonne input_ids", placeholder="Question")
            colonne_attention_mask = gr.Textbox(label="Colonne attention_mask", placeholder="Reponse")
            colonne_labels = gr.Textbox(label="Colonne labels (si applicable)")

        with gr.Row():
            load_button = gr.Button("Charger Le DataSet")
        with gr.Row():
            load_output = gr.Textbox(label="État du chargement du DataSet")

        with gr.Row():
            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")
            stop_fine_tune_button = gr.Button("Stopper Le Fine-Tuning")  # Ajout du bouton STOP

        with gr.Row():
            fine_tune_output = gr.Textbox(label="État du Fine-Tuning")
            logs_output = gr.Textbox(label="Logs d'entraînement", interactive=False)

        def load_dataset(dataset_path, dataset_file, split, colonne_input_ids, colonne_attention_mask, colonne_labels):
            try:
                response = requests.get(f"{dataset_path}/resolve/main/{dataset_file}")
                dataset_lines = response.text.strip().split('\n')
                dataset_dict = [json.loads(line) for line in dataset_lines if line]
                return "DataSet chargé avec succès"
            except Exception as e:
                return str(e)

        load_button.click(
            load_dataset,
            inputs=[dataset_path, dataset_file, split, colonne_input_ids, colonne_attention_mask, colonne_labels],
            outputs=load_output
        )

        fine_tune_button.click(
            lambda *args: (threading.Thread(target=fine_tune_model, args=args, daemon=True).start(), "Entraînement lancé."),
            inputs=[dataset_path, dataset_file, split, colonne_input_ids, colonne_attention_mask, colonne_labels, epochs, batch_size, num_cpus, ram_size, cpu_limit, ram_limit, gpu_limit],  # Ajustement du nombre d'arguments
            outputs=fine_tune_output
        )

        stop_fine_tune_button.click(
            stop_fine_tune,
            None,
            fine_tune_output
        )

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