Spaces:
Sleeping
Sleeping
File size: 1,918 Bytes
ba0e6c9 cc825ac ba0e6c9 |
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 |
import gradio as gr
import torch
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
# On charge le pipeline une fois au démarrage
pipe = load_model()
# 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"]
# Interface Gradio
with gr.Blocks() as demo:
gr.Markdown("# 🐟 BRVM Finetuner (Mistral-7B)")
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)
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)
demo.launch()
|