Update app.py

app.py

@@ -4,32 +4,40 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
 
 MODEL_DIR = "./mistral-7b-brvm-finetuned"
+FULL_MODEL_DIR = "./mistral-7b-brvm-full-finetuned"
 
-# Fonction d’entraînement (appelle ton script de fine-tuning)
+# Fonction d’entraînement LoRA
 def train_model():
-    os.system("python finetune.py")  # tu mets ton code d'entraînement dans finetune.py
-    return "✅ Entraînement terminé ! Le modèle est sauvegardé dans " + MODEL_DIR
+    os.system("python finetune.py")
+    return "✅ Entraînement LoRA terminé ! Modèle sauvegardé dans " + MODEL_DIR
 
-# Chargement du modèle (fine-tuné si dispo, sinon base)
+# Fonction d’entraînement Full Fine-tune
+def train_model_full():
+    os.system("python finetune_full.py")
+    return "✅ Full fine-tune terminé ! Modèle sauvegardé dans " + FULL_MODEL_DIR
+
+# Chargement du modèle (par défaut LoRA, sinon base)
 def load_model():
-    model_name = MODEL_DIR if os.path.exists(MODEL_DIR) else "mistralai/Mistral-7B-Instruct-v0.3"
+    if os.path.exists(FULL_MODEL_DIR):
+        model_name = FULL_MODEL_DIR
+    elif os.path.exists(MODEL_DIR):
+        model_name = MODEL_DIR
+    else:
+        model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map="auto",  # Accelerate gère la répartition CPU/GPU
+        device_map="auto",
         torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         trust_remote_code=True,
     )
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-    # Ne PAS passer device quand on utilise accelerate
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-    return pipe
-
-# On charge le pipeline une fois au démarrage
+# Charger pipeline
 pipe = load_model()
 
-# Fonction de test du modèle
+# Fonction de test
 def chat(prompt):
     outputs = pipe(prompt, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.9)
     return outputs[0]["generated_text"]
@@ -39,9 +47,12 @@ with gr.Blocks() as demo:
     gr.Markdown("# 🐟 BRVM Finetuner (Mistral-7B)")
 
     with gr.Tab("🚀 Entraînement"):
-        train_btn = gr.Button("Lancer l’entraînement")
+        with gr.Row():
+            train_btn = gr.Button("Lancer LoRA Fine-tuning")
+            train_full_btn = gr.Button("Lancer Full Fine-tuning")
         train_output = gr.Textbox(label="Logs")
         train_btn.click(fn=train_model, outputs=train_output)
+        train_full_btn.click(fn=train_model_full, outputs=train_output)
 
     with gr.Tab("💬 Tester le modèle"):
         input_text = gr.Textbox(label="Votre question :", placeholder="Posez une question...")
@@ -49,4 +60,4 @@ with gr.Blocks() as demo:
         submit_btn = gr.Button("Envoyer")
         submit_btn.click(fn=chat, inputs=input_text, outputs=output_text)
 
-demo.launch()
+demo.launch()