Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, Mistral3ForConditionalGeneration, TextIteratorStreamer
 from threading import Thread
 import re
 import time
@@ -47,13 +47,13 @@ class SinaReasonMedicalChat:
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
-            self.model = AutoModelForCausalLM.from_pretrained(
+            self.model = Mistral3ForConditionalGeneration.from_pretrained(
                 MODEL_NAME,
                 dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
-                #device_map="auto" if DEVICE == "cuda" else None,
-                trust_remote_code=True,
-                low_cpu_mem_usage=True
-            )
+                device_map="auto" if DEVICE == "cuda" else None,
+                #trust_remote_code=True,
+                #low_cpu_mem_usage=True
+            ).eval()
             
 
                 
@@ -78,7 +78,7 @@ class SinaReasonMedicalChat:
         
         return thinking, response
 
-    @spaces.GPU
+    @spaces.GPU(duration=120)
     def medical_chat_stream(self, message: str, history: List[List[str]], max_tokens: int = 1024, 
                           temperature: float = 0.7, top_p: float = 0.95) -> Iterator[Tuple[str, List[List[str]]]]:
         """Stream medical reasoning responses with thinking display"""