Update app.py

app.py

@@ -1,9 +1,18 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer      
-import gradio as gr      
-import torch      
-    
+import gradio as gr           
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+base_model_id = "mistralai/Mistral-7B-v0.1"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config)
 base_model = AutoModelForCausalLM.from_pretrained(
-    Mistral,  # Mistral, same as before
+    base_model_id,  # Mistral, same as before
     quantization_config=bnb_config,  # Same quantization config as before
     device_map="auto",
     trust_remote_code=True,