File size: 4,832 Bytes
eb96a2e
 
 
 
 
 
 
 
 
 
e39c621
 
eb96a2e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e39c621
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eb96a2e
e39c621
 
 
eb96a2e
 
e39c621
eb96a2e
 
 
 
 
e39c621
 
eb96a2e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#BS_app.py_02
#Training NOK

#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 le fine-tuning
#def fine_tune_model(dataset_path, dataset_file, epochs, batch_size, prefix):
#    # 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})
#
#    # 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)




def fine_tune_model(dataset_path, dataset_file, epochs, batch_size, prefix):
    # Récupération du fichier à partir de l'URL fournie
    response = requests.get(dataset_path)
    dataset = list(response.iter_lines())

    # Préparation des données
    dataset = Dataset.from_dict({'data': dataset})
    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=f"./{prefix}_{model_name.split('/')[-1]}",
        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=f"{prefix}_{model_name.split('/')[-1]}",
        hub_strategy="checkpoint",
        hub_token=hf_token,
    )

    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 avec un préfixe
    trainer.save_model(f"./{prefix}_{model_name.split('/')[-1]}")
    tokenizer.save_pretrained(f"./{prefix}_{model_name.split('/')[-1]}")

    # Push vers Hugging Face Hub
    api = HfApi()
    api.upload_folder(
        folder_path=f"./{prefix}_{model_name.split('/')[-1]}",
        repo_id=f"{prefix}_{model_name.split('/')[-1]}",
        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")
            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)
            prefix = gr.Textbox(label="Préfixe pour les fichiers sauvegardés")

            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, dataset_file, epochs, batch_size, prefix],
            outputs=fine_tune_output
        )

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