Update app.py

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 import spaces 
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from megatron import TransformerModel
 
 
 
 tokenizer = AutoTokenizer.from_pretrained("nvidia/mamba2-hybrid-8b-3t-128k")
-model = AutoModelForCausalLM.from_pretrained("nvidia/mamba2-hybrid-8b-3t-128k", device_map="auto")  # to("cuda:0") 
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
+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