File size: 2,219 Bytes
6d581b7 |
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 |
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
# Remplacez par le modèle de base et l'adaptateur LoRA que vous utilisez
BASE_MODEL = "bigcode/starcoder2-3b"
ADAPTER_REPO = "simnJS/autotrain-fxp6j-p5s8i"
# 1. Charger le tokenizer
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
# 2. Charger le modèle de base en FP16 sur le GPU (si vous avez la VRAM nécessaire)
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=torch.float16,
device_map="auto" # Permet de placer le modèle sur le GPU automatiquement
)
# 3. Charger l'adaptateur LoRA
model = PeftModel.from_pretrained(
base_model,
ADAPTER_REPO,
torch_dtype=torch.float16
)
# 4. Fonction pour générer une réponse
def generate_answer(user_message, history):
"""
user_message: le dernier message de l'utilisateur
history: liste de tuples (message_utilisateur, réponse_modèle)
"""
# Construire le prompt en tenant compte de l'historique si besoin
# Ici, on fait simple et on utilise juste le dernier message
prompt = user_message
# Encoder le prompt
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
# Générer la réponse
outputs = model.generate(
**inputs,
max_new_tokens=100,
temperature=0.7,
do_sample=True,
top_p=0.9
)
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Ajouter dans l'historique
history.append((user_message, answer))
return history, history
# 5. Interface Gradio de type chatbot
with gr.Blocks() as demo:
gr.Markdown("# Chat avec mon modèle LoRA Verse")
chatbot = gr.Chatbot()
msg = gr.Textbox(label="Tapez votre message ici...")
state = gr.State([]) # pour stocker l'historique des messages
def submit_message(user_message, history):
return generate_answer(user_message, history)
msg.submit(submit_message, inputs=[msg, state], outputs=[chatbot, state])
# ou un bouton si vous préférez
# send_btn = gr.Button("Envoyer")
# send_btn.click(fn=submit_message, inputs=[msg, state], outputs=[chatbot, state])
demo.launch()
|