Update app.py

app.py

@@ -1,212 +1,98 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
-from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Configuration du modèle
-ADAPTER_MODEL_NAME = "analist/llama3.1-8B-omnimed-rl"
-BASE_MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B" # Modèle de base pour Llama 3.1 8B
-DEFAULT_SYSTEM_PROMPT = """Vous êtes OmniMed, un assistant médical IA conçu pour aider les professionnels de santé dans leurs tâches quotidiennes.
-Répondez de manière précise, concise et professionnelle aux questions médicales.
-Basez vos réponses sur des connaissances médicales établies et indiquez clairement lorsque vous n'êtes pas certain d'une information."""
+# Définir le modèle et le tokenizer
+# Utilisation d'un modèle français pour le domaine médical
+MODEL_NAME = "mistralai/Mistral-7B-v0.1"  # Vous pouvez utiliser un modèle plus adapté au français comme "camembert" ou un modèle médical spécifique
 
-# Définition des paramètres
-MAX_NEW_TOKENS = 1024
-TEMPERATURE = 0.7
-TOP_P = 0.9
-REPETITION_PENALTY = 1.1
-
-# Configuration pour la quantification 4-bit
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
-)
+# Fonction pour charger le modèle et le tokenizer
+def load_model():
+    print("Chargement du modèle et du tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        load_in_8bit=True  # Quantification pour réduire l'utilisation de la mémoire
+    )
+    return model, tokenizer
 
-# Chargement du modèle et du tokenizer
-print("Chargement du modèle de base et du tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(ADAPTER_MODEL_NAME)
+# Charger le modèle et le tokenizer
+model, tokenizer = load_model()
 
-print("Chargement du modèle de base quantifié...")
-base_model = AutoModelForCausalLM.from_pretrained(
-    ADAPTER_MODEL_NAME,
-    quantization_config=bnb_config,
-    device_map="auto",
-    trust_remote_code=True
+# Créer un pipeline de génération de texte
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_length=512,
+    do_sample=True,
+    temperature=0.7,
+    top_p=0.9,
 )
 
-print("Application des adaptateurs...")
-model = PeftModel.from_pretrained(
-    base_model,
-    device_map="auto",
-)
-
-print("Modèle et tokenizer chargés avec succès!")
+# Message initial pour orienter le modèle vers le domaine médical
+SYSTEM_PROMPT = """Tu es un assistant médical IA nommé MediBot. Tu fournis des informations médicales générales et des conseils de santé basés sur des données scientifiques.
+Important: Tu n'es pas un médecin et tu dois toujours recommander à l'utilisateur de consulter un professionnel de santé pour un diagnostic ou un traitement spécifique.
+Réponds aux questions médicales de manière précise et claire, en te basant sur des informations médicales vérifiées."""
 
 # Fonction pour générer une réponse
-def generate_response(message, chat_history, system_prompt, temperature=TEMPERATURE, max_tokens=MAX_NEW_TOKENS):
-    # Construction du contexte de chat
-    chat_context = []
-    
-    # Ajout du prompt système s'il existe
-    if system_prompt.strip():
-        chat_context.append({"role": "system", "content": system_prompt})
-    else:
-        chat_context.append({"role": "system", "content": DEFAULT_SYSTEM_PROMPT})
-    
-    # Ajout de l'historique des conversations
-    for user_msg, assistant_msg in chat_history:
-        chat_context.append({"role": "user", "content": user_msg})
-        chat_context.append({"role": "assistant", "content": assistant_msg})
-    
-    # Ajout du message actuel
-    chat_context.append({"role": "user", "content": message})
-    
-    # Préparation du texte d'entrée avec le template de chat
-    input_text = tokenizer.apply_chat_template(
-        chat_context,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    
-    # Tokenisation de l'entrée
-    model_inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
-    
-    # Mise à jour des paramètres de génération
-    model.generation_config.temperature = temperature
-    model.generation_config.max_new_tokens = max_tokens
-    
-    # Génération de la réponse
-    with torch.no_grad():
-        generated_ids = model.generate(
-            **model_inputs,
-            use_cache=True,
-        )
+def generate_response(message, chat_history):
+    # Construire le prompt avec l'historique de conversation et le nouveau message
+    prompt = SYSTEM_PROMPT + "\n\n"
     
-    # Extraction uniquement de la nouvelle partie générée
-    new_tokens = generated_ids[0][model_inputs.input_ids.shape[1]:]
-    response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+    # Ajouter l'historique du chat
+    for user_msg, bot_msg in chat_history:
+        prompt += f"Utilisateur: {user_msg}\nMediBot: {bot_msg}\n\n"
     
-    return response
+    # Ajouter le nouveau message
+    prompt += f"Utilisateur: {message}\nMediBot:"
     
-    # Construction du contexte de chat
-    chat_context = []
+    # Générer la réponse
+    response = generator(prompt, max_new_tokens=256)[0]["generated_text"]
     
-    # Ajout du prompt système s'il existe
-    if system_prompt.strip():
-        chat_context.append({"role": "system", "content": system_prompt})
-    else:
-        chat_context.append({"role": "system", "content": DEFAULT_SYSTEM_PROMPT})
+    # Extraire seulement la partie réponse du modèle
+    response_only = response.split("MediBot:")[-1].strip()
     
-    # Ajout de l'historique des conversations
-    for user_msg, assistant_msg in chat_history:
-        chat_context.append({"role": "user", "content": user_msg})
-        chat_context.append({"role": "assistant", "content": assistant_msg})
+    # Ajouter un rappel de consulter un professionnel
+    if len(response_only) > 0 and not "consulter un professionnel" in response_only.lower():
+        response_only += "\n\n(N'oubliez pas que ces informations sont générales. Pour des conseils personnalisés, consultez un professionnel de santé.)"
     
-    # Ajout du message actuel
-    chat_context.append({"role": "user", "content": message})
-    
-    # Tokenisation et génération
-    input_ids = tokenizer.apply_chat_template(
-        chat_context,
-        return_tensors="pt"
-    ).to(model.device)
-    
-    # Génération de la réponse
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids=input_ids,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=TOP_P,
-            repetition_penalty=REPETITION_PENALTY,
-            do_sample=temperature > 0,
-        )
-    
-    # Décodage de la réponse
-    generated_text = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
-    return generated_text
+    return response_only
 
-# Interface utilisateur Gradio
-with gr.Blocks(css="footer {visibility: hidden}") as demo:
-    gr.Markdown(
-        """
-        # OmniMed - Assistant médical IA
-        
-        Ce chatbot alimenté par le modèle Llama 3.1 fine-tuné est conçu pour aider les professionnels de santé dans leurs tâches quotidiennes.
-        Posez des questions médicales et recevez des réponses précises et professionnelles.
-        """
-    )
-    
-    with gr.Row():
-        with gr.Column(scale=4):
-            chatbot = gr.Chatbot(height=500, show_copy_button=True)
-            
-            with gr.Row():
-                msg = gr.Textbox(
-                    placeholder="Posez votre question médicale ici...",
-                    show_label=False,
-                    container=False,
-                    scale=12
-                )
-                submit_btn = gr.Button("Envoyer", scale=1)
-            
-            with gr.Row():
-                clear_btn = gr.Button("Effacer la conversation")
-        
-        with gr.Column(scale=1):
-            with gr.Accordion("Paramètres avancés", open=False):
-                system_prompt = gr.Textbox(
-                    label="Prompt système",
-                    placeholder="Instructions pour l'assistant...",
-                    value=DEFAULT_SYSTEM_PROMPT,
-                    lines=5
-                )
-                temperature_slider = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=TEMPERATURE,
-                    step=0.05,
-                    label="Température"
-                )
-                max_tokens_slider = gr.Slider(
-                    minimum=128,
-                    maximum=2048,
-                    value=MAX_NEW_TOKENS,
-                    step=64,
-                    label="Tokens max."
-                )
-    
-    # Logique des événements
-    def respond(message, chat_history, system_prompt, temperature, max_tokens):
-        if not message.strip():
-            return "", chat_history
-        
-        response = generate_response(message, chat_history, system_prompt, temperature, max_tokens)
-        chat_history.append((message, response))
-        return "", chat_history
+# Fonction pour gérer l'historique du chat Gradio
+def chatbot_response(message, history):
+    bot_message = generate_response(message, history)
+    history.append((message, bot_message))
+    return "", history
+
+# Interface Gradio
+with gr.Blocks(title="MediBot - Assistant Médical IA") as demo:
+    gr.Markdown("# MediBot - Votre Assistant Médical basé sur l'IA")
+    gr.Markdown("""
+    ### ⚠️ IMPORTANT - AVERTISSEMENT MÉDICAL ⚠️
+    Ce chatbot utilise l'intelligence artificielle pour fournir des informations médicales générales.
+    Il ne remplace en aucun cas l'avis d'un professionnel de santé qualifié.
+    Pour toute question médicale sérieuse, veuillez consulter un médecin.
+    """)
     
-    msg.submit(
-        respond,
-        [msg, chatbot, system_prompt, temperature_slider, max_tokens_slider],
-        [msg, chatbot]
-    )
+    chatbot = gr.Chatbot(height=500)
+    msg = gr.Textbox(placeholder="Posez votre question médicale ici...", label="Votre question")
+    clear = gr.Button("Effacer la conversation")
     
-    submit_btn.click(
-        respond,
-        [msg, chatbot, system_prompt, temperature_slider, max_tokens_slider],
-        [msg, chatbot]
-    )
+    # Configurer les événements
+    msg.submit(chatbot_response, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
     
-    clear_btn.click(lambda: [], None, chatbot)
-
     gr.Markdown("""
-    ### Note importante:
-    Ce chatbot est destiné uniquement à l'assistance des professionnels de santé et ne remplace pas l'avis médical d'un spécialiste.
-    Toutes les informations fournies doivent être vérifiées par un professionnel qualifié.
+    ### Exemples de questions que vous pouvez poser:
+    - Quels sont les symptômes courants de la grippe?
+    - Comment prévenir l'hypertension artérielle?
+    - Quels aliments sont recommandés pour les diabétiques?
+    - Quelles sont les causes fréquentes des migraines?
     """)
 
-# Lancement de l'application
+# Lancer l'application
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.launch(share=True)  # share=True permet de générer un lien public temporaire