Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
 import spaces 
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from megatron import TransformerModel
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
 
-tokenizer = AutoTokenizer.from_pretrained("nvidia/mamba2-hybrid-8b-3t-128k")
-model = TransformerModel.from_pretrained("nvidia/mamba2-hybrid-8b-3t-128k", device_map="auto")  # to("cuda:0") 
-
 
 
 # Optionnel : Fixer une graine aléatoire pour la reproductibilit