#testing bloom1b training

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

# 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 le fine-tuning
def fine_tune_model(dataset_path, epochs, batch_size):
    # Chargement du dataset
    if dataset_path.startswith("huggingface://"):
        dataset = load_dataset(dataset_path.replace("huggingface://", ""))
    else:
        dataset = load_dataset('text', data_files={'train': dataset_path})

    # Préparation des données
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # Configuration de l'entraînement
    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        save_steps=10_000,
        save_total_limit=2,
        push_to_hub=True,
        hub_model_id=model_name,
        hub_strategy="checkpoint",
        hub_token=HfFolder.get_token(),
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset['train'],
    )

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

    # Sauvegarde du modèle avec un préfixe
    trainer.save_model(f"./FT01_{model_name.split('/')[-1]}")
    tokenizer.save_pretrained(f"./FT01_{model_name.split('/')[-1]}")

    # Push vers Hugging Face Hub
    api = HfApi()
    api.upload_folder(
        folder_path=f"./FT01_{model_name.split('/')[-1]}",
        repo_id=model_name,
        repo_type="model"
    )

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

# 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="Entrez votre message pour obtenir une réponse du modèle"
        )

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

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

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

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