Update app.py

app.py

@@ -1,10 +1,11 @@
 #testing bloom1b training
 
+
 import gradio as gr
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
 from datasets import load_dataset
-from huggingface_hub import HfFolder
+from huggingface_hub import HfApi, HfFolder
 
 # Récupérer le token depuis les variables d'environnement
 hf_token = os.getenv("MisterAI_bigscience_bloom_560m")
@@ -28,23 +29,15 @@ def generate_response(input_text):
     return response
 
 # Fonction pour le fine-tuning
-def fine_tune_model(dataset_path, epochs, batch_size, input_columns, output_column):
+def fine_tune_model(dataset_path, epochs, batch_size):
     # Chargement du dataset
     if dataset_path.startswith("https://huggingface.co/datasets/"):
-        dataset = load_dataset(dataset_path.replace("https://huggingface.co/datasets/", ""))
+        dataset = load_dataset('json', data_files={'train': dataset_path})
     else:
         dataset = load_dataset('json', data_files={'train': dataset_path})
 
     # Préparation des données
-    def preprocess_function(examples):
-        inputs = [f"{input_columns[0]}: {examples[input_columns[0]][i]}\n{output_column}: {examples[output_column][i]}" for i in range(len(examples[input_columns[0]]))]
-        model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
-        return model_inputs
-
-    tokenized_dataset = dataset['train'].map(preprocess_function, batched=True)
-
-    # Configuration du Data Collator
-    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
     # Configuration de l'entraînement
     training_args = TrainingArguments(
@@ -63,12 +56,24 @@ def fine_tune_model(dataset_path, epochs, batch_size, input_columns, output_colu
         model=model,
         args=training_args,
         data_collator=data_collator,
-        train_dataset=tokenized_dataset,
+        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
@@ -87,15 +92,13 @@ with gr.Blocks() as demo:
             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)
-            input_columns = gr.Textbox(label="Colonnes d'entrée (séparées par des virgules)", value="question")
-            output_column = gr.Textbox(label="Colonne de sortie", value="chosen")
             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, input_columns, output_column],
+            inputs=[dataset_path, epochs, batch_size],
             outputs=fine_tune_output
         )
 
@@ -103,3 +106,6 @@ with gr.Blocks() as demo:
 if __name__ == "__main__":
     demo.launch()
 
+
+
+