Update app.py

app.py

@@ -1,9 +1,11 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
+from peft import PeftModel, PeftConfig
 
 # Configuration du modèle
-MODEL_NAME = "analist/llama3.1-8B-omnimed-rl"
+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."""
@@ -14,15 +16,33 @@ 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,
+)
+
 # Chargement du modèle et du tokenizer
-print("Chargement du modèle et du tokenizer...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
+print("Chargement du modèle de base et du tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
+
+print("Chargement du modèle de base quantifié...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_NAME,
+    quantization_config=bnb_config,
     device_map="auto",
     trust_remote_code=True
 )
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+print("Application des adaptateurs...")
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER_MODEL_NAME,
+    device_map="auto",
+)
+
 print("Modèle et tokenizer chargés avec succès!")
 
 # Fonction pour générer une réponse