Update app.py

app.py

@@ -1,148 +1,49 @@
+import gradio as gr
 import torch
-from datasets import load_dataset
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig,
-    TrainingArguments,
-)
-from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-from trl import SFTTrainer
-
-# 1. Configurations
-base_model = "mistralai/Mistral-7B-Instruct-v0.3"
-new_model_dir = "./mistral-7b-brvm-finetuned"
-output_dir = "./results"
-
-# 2. Device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Utilisation du périphérique: {device}")
-if torch.cuda.is_available():
-    print(f"GPU: {torch.cuda.get_device_name(0)} - "
-          f"Mémoire: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB")
-
-# 3. Dataset
-dataset = load_dataset("lamekemal/brvm_finetune")
-
-# 4. Quantization
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=False,
-)
-
-# 5. Charger modèle + tokenizer
-model = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    quantization_config=bnb_config,
-    device_map="auto",
-    trust_remote_code=True,
-)
-model.config.use_cache = False
-model = prepare_model_for_kbit_training(model)
-
-tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"
-
-# ============================================
-# 10. Preprocessing (max_seq_length=512)
-# ============================================
-def tokenize_function(examples):
-    texts = [
-        f"Instruction: {instr}\nRéponse: {resp}"
-        for instr, resp in zip(examples["instruction"], examples["response"])
-    ]
-    return tokenizer(
-        texts,
-        truncation=True,
-        padding="max_length",
-        max_length=512,
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import os
+
+MODEL_DIR = "./mistral-7b-brvm-finetuned"
+
+# Fonction d’entraînement (appelle ton script de fine-tuning)
+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
+
+# Chargement du modèle (fine-tuné si dispo, sinon base)
+def load_model():
+    model_name = MODEL_DIR if os.path.exists(MODEL_DIR) else "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",
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        trust_remote_code=True,
     )
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+    return pipe
 
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
-# 6. LoRA config
-lora_config = LoraConfig(
-    r=16,
-    lora_alpha=32,
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
-                    "gate_proj", "up_proj", "down_proj"],
-    lora_dropout=0.05,
-    bias="none",
-    task_type="CAUSAL_LM",
-)
-model = get_peft_model(model, lora_config)
-
-sft_config = SFTConfig(
-    output_dir=output_dir,
-    num_train_epochs=3,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=4,
-    optim="paged_adamw_32bit",
-    save_steps=100,
-    logging_steps=10,
-    learning_rate=2e-4,
-    fp16=False,
-    bf16=torch.cuda.is_available(),
-    max_grad_norm=0.3,
-    warmup_ratio=0.03,
-    group_by_length=True,
-    lr_scheduler_type="cosine",
-    report_to="tensorboard",
-    evaluation_strategy="steps",
-    eval_steps=100,
-    save_total_limit=2,
-    load_best_model_at_end=True,
-    metric_for_best_model="eval_loss",
-    max_seq_length=512,
-    packing=False,
-)
-
-# ============================================
-# 12. TrainingArguments
-# ============================================
-use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+# On charge le pipeline une fois au démarrage
+pipe = load_model()
 
-training_args = TrainingArguments(
-    output_dir=output_dir,
-    num_train_epochs=3,
-    per_device_train_batch_size=4,
-    gradient_accumulation_steps=2,
-    optim="paged_adamw_8bit",
-    save_steps=100,
-    logging_steps=10,
-    learning_rate=2e-4,
-    #fp16=not use_bf16,
-    bf16=True,
-    max_grad_norm=0.3,
-    warmup_ratio=0.03,
-    group_by_length=True,
-    lr_scheduler_type="cosine",
-    report_to="tensorboard",
-    eval_strategy="steps",   # <-- corrige le nom
-    eval_steps=100,
-    save_total_limit=2,
-    load_best_model_at_end=True,
-    metric_for_best_model="eval_loss",
-)
+# Fonction de test du modèle
+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"]
 
-# ============================================
-# 13. Trainer
-# ============================================
-trainer = SFTTrainer(
-    model=model,
-    train_dataset=tokenized_datasets["train"],
-    eval_dataset=tokenized_datasets["validation"],
-    peft_config=lora_config,
-    args=training_args
-)
+# Interface Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# 🐟 BRVM Finetuner (Mistral-7B)")
 
-# 9. Fine-tuning
-trainer.train()
+    with gr.Tab("🚀 Entraînement"):
+        train_btn = gr.Button("Lancer l’entraînement")
+        train_output = gr.Textbox(label="Logs")
+        train_btn.click(fn=train_model, outputs=train_output)
 
-# 10. Sauvegarde locale
-trainer.save_model(new_model_dir)
+    with gr.Tab("💬 Tester le modèle"):
+        input_text = gr.Textbox(label="Votre question :", placeholder="Posez une question...")
+        output_text = gr.Textbox(label="Réponse du modèle")
+        submit_btn = gr.Button("Envoyer")
+        submit_btn.click(fn=chat, inputs=input_text, outputs=output_text)
 
-print(f"✅ Modèle LoRA sauvegardé localement dans {new_model_dir}")
+demo.launch()