Update app.py

app.py

@@ -5,6 +5,7 @@ from typing import Iterator
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import bitsandbytes
 
 MAX_MAX_NEW_TOKENS = 4096
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -16,7 +17,7 @@ DESCRIPTION = """\
 
 # Load model with appropriate device configuration
 def load_model():
-    model_id = "CreitinGameplays/Llama-3.1-8B-R1-v0.1"
+    model_id = "CreitinGameplays/Mistral-Nemo-12B-R1-v0.1"
     device = "cuda" if torch.cuda.is_available() else "cpu"
     
     # If using CPU, load in 32-bit to avoid potential issues with 16-bit operations
@@ -30,7 +31,8 @@ def load_model():
         model = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
-            device_map="auto"
+            device_map="auto",
+            load_in_8bit=True
         )
     tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side='left')
     tokenizer.use_default_system_prompt = False