Train / run_qwen_model.py
kvn420's picture
Create run_qwen_model.py
4c516c0 verified
# app.py
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import __version__ as transformers_version
MODEL_NAME = "kvn420/Tenro_V4.1"
RECOMMENDED_TRANSFORMERS_VERSION = "4.37.0"
print(f"Version de Transformers : {transformers_version}")
if transformers_version < RECOMMENDED_TRANSFORMERS_VERSION:
print(f"Attention : Version Transformers ({transformers_version}) < Recommandée ({RECOMMENDED_TRANSFORMERS_VERSION}). Mettez à jour.")
# --- Chargement du modèle et du tokenizer (une seule fois) ---
print(f"Chargement du tokenizer pour : {MODEL_NAME}")
try:
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True
)
print("Tokenizer chargé.")
print(f"Chargement du modèle : {MODEL_NAME}")
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16, # ou "auto"
trust_remote_code=True,
device_map="auto" # Utilise le GPU si disponible, sinon CPU
)
print(f"Modèle chargé sur {model.device}.")
# Définir pad_token_id si manquant (important pour la génération)
if tokenizer.pad_token_id is None:
tokenizer.pad_token_id = tokenizer.eos_token_id
print(f"tokenizer.pad_token_id défini sur eos_token_id: {tokenizer.eos_token_id}")
except Exception as e:
print(f"Erreur critique lors du chargement du modèle/tokenizer : {e}")
# Lever l'erreur pour que Gradio l'affiche ou la loggue
raise gr.Error(f"Impossible de charger le modèle ou le tokenizer: {e}. Vérifiez les logs du Space.")
# --- Fin du chargement ---
def chat_interaction(user_input, history):
"""
Fonction appelée par Gradio pour chaque interaction de chat.
history est une liste de paires [user_message, assistant_message]
"""
if model is None or tokenizer is None:
return "Erreur: Modèle ou tokenizer non initialisé."
# Construire le prompt avec l'historique pour le modèle
messages_for_template = []
# Ajouter un message système par défaut si l'historique est vide et le premier message n'est pas un système
# Ou si vous voulez toujours un message système spécifique.
# Note: Le chat_template de Qwen ajoute déjà un message système par défaut.
# Adaptez selon le comportement exact souhaité.
# messages_for_template.append({"role": "system", "content": "Tu es Qwen, un assistant IA serviable."})
for user_msg, assistant_msg in history:
messages_for_template.append({"role": "user", "content": user_msg})
messages_for_template.append({"role": "assistant", "content": assistant_msg})
messages_for_template.append({"role": "user", "content": user_input})
try:
prompt_tokenized = tokenizer.apply_chat_template(
messages_for_template,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt"
).to(model.device)
outputs = model.generate(
prompt_tokenized,
max_new_tokens=512, # Augmenté pour des réponses potentiellement plus longues
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.pad_token_id
)
response_text = tokenizer.decode(outputs[0][prompt_tokenized.shape[-1]:], skip_special_tokens=True)
# Nettoyage simple (optionnel, dépend du modèle)
response_text = response_text.replace("<|im_end|>", "").strip()
if response_text.startswith("assistant\n"): # Parfois Qwen ajoute cela
response_text = response_text.split("assistant\n", 1)[-1].strip()
return response_text
except Exception as e:
print(f"Erreur pendant la génération : {e}")
return f"Désolé, une erreur est survenue : {e}"
# Création de l'interface Gradio
# Utilisation de `gr.ChatInterface` qui gère l'historique automatiquement.
iface = gr.ChatInterface(
fn=chat_interaction,
title=f"Chat avec {MODEL_NAME}",
description=f"Interface de démonstration pour le modèle {MODEL_NAME}. Le modèle est hébergé sur Hugging Face et chargé ici.",
chatbot=gr.Chatbot(height=600),
textbox=gr.Textbox(placeholder="Posez votre question ici...", container=False, scale=7),
retry_btn="Réessayer",
undo_btn="Annuler",
clear_btn="Effacer la conversation",
submit_btn="Envoyer"
)
# Lancer l'application (pour un test local, ou si vous n'utilisez pas `if __name__ == "__main__":` dans Spaces)
# iface.launch() # Décommentez pour un test local facile
# Pour Spaces, il est souvent préférable de laisser Spaces gérer le lancement
# si vous utilisez le SDK Gradio directement dans la configuration du Space.
# Si vous exécutez ce script directement avec python app.py, il faut iface.launch().
# Dans le contexte d'un Space Gradio, le iface est généralement ce qui est "retourné" implicitement.
# Si vous voulez contrôler le lancement (ex: pour des options), utilisez :
# if __name__ == "__main__":
# iface.launch()
# Mais pour un Space Gradio simple, juste définir `iface` peut suffire.
# La convention est de lancer si le script est exécuté directement.
if __name__ == '__main__':
iface.launch() # Permet de tester localement `python app.py`