Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, Mistral3ForConditionalGeneration, TextIteratorStreamer
+from transformers import AutoTokenizer, Mistral3ForConditionalGeneration, TextIteratorStreamer,BitsAndBytesConfig
 from threading import Thread
 import re
 import time
@@ -41,16 +41,23 @@ class SinaReasonMedicalChat:
         try:
             print(f"Loading medical model: {MODEL_NAME}")
             self.tokenizer = AutoTokenizer.from_pretrained(
-                "mistralai/Magistral-Small-2509"
+                "mistralai/Magistral-Small-2509",
+                tokenizer_type="mistral"
             )
-            
+            bnb_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",         # normal float4
+                bnb_4bit_compute_dtype=torch.bfloat16   # computation precision
+            )
+                        
             # Add padding token if not present
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
             self.model = Mistral3ForConditionalGeneration.from_pretrained(
                 MODEL_NAME,
-                dtype=torch.bfloat16
+                quantization_config=bnb_config
             )