chat-medical / app.py
Ronaldodev's picture
Update app.py
3f982f9 verified
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import gradio as gr
import re
from pathlib import Path
# ====== CONFIGURATION ======
MODEL_NAME = "microsoft/phi-3-mini-4k-instruct"
SYSTEM_PROMPT_FILE = "infos_medical.txt" # ton fichier .txt
# ====== CHARGEMENT DU MODÈLE ET TOKENIZER ======
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map="auto"
)
print(f"✅ Modèle '{MODEL_NAME}' chargé avec succès.")
# ====== CHARGEMENT DU CONTEXTE MÉDICAL ======
try:
context_text = Path(SYSTEM_PROMPT_FILE).read_text(encoding="utf-8").strip()
except FileNotFoundError:
context_text = "Aucune donnée médicale fournie."
print("⚠️ Fichier infos_medical.txt introuvable — le chatbot fonctionnera sans contexte.")
# Nettoyage du texte
context_text = context_text.replace("\n", " ")
context_text = re.sub(r"\s+", " ", context_text)
context_text = context_text[:3000] # limite la longueur pour garder de la mémoire
# ====== INSTRUCTION SYSTÈME ======
system_instruction = (
"Tu es un assistant médical virtuel. "
"Tu réponds en français, poliment et de façon factuelle. "
"Tu fournis uniquement des informations générales de santé. "
"Tu ne poses pas de diagnostic et ne prescris pas de traitement. "
"Si la demande est trop spécifique, réponds : "
"\"Je ne sais pas — consultez un professionnel de santé.\""
)
# ====== CONSTRUCTION DU PROMPT ======
def build_prompt(history, user_message):
"""
Construit un prompt propre : instruction + contexte + historique + message utilisateur.
"""
lines = [
f"System: {system_instruction}",
"",
f"Contexte médical de référence : {context_text}",
"",
]
for user, bot in history:
lines.append(f"User: {user}")
if bot:
lines.append(f"Bot: {bot}")
lines.append(f"User: {user_message}")
lines.append("Bot:")
return "\n".join(lines)
# ====== FONCTION DE RÉPONSE ======
def chat(user_message, history):
prompt = build_prompt(history, user_message)
inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=300,
temperature=0.6,
top_p=0.9,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Récupération du texte après le dernier "Bot:"
if "Bot:" in reply:
reply = reply.split("Bot:")[-1].strip()
history.append((user_message, reply))
return history, history
# ====== INTERFACE GRADIO ======
with gr.Blocks(theme=gr.themes.Soft(primary_hue="green")) as demo:
gr.Markdown("## 🩺 Chatbot Médical — Assistant virtuel de santé")
gr.Markdown(f"📘 **Contexte chargé :** {len(context_text.split())} mots de référence médicale.")
chatbot = gr.Chatbot(height=400)
msg = gr.Textbox(label="Posez une question médicale…")
clear = gr.Button("🧹 Effacer la conversation")
msg.submit(chat, [msg, chatbot], [chatbot, chatbot])
clear.click(lambda: None, None, chatbot, queue=False)
demo.launch()