Update app.py

app.py

@@ -2,14 +2,24 @@ import os
 import bitsandbytes as bnb
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 access_token = os.environ["GATED_ACCESS_TOKEN"]
 
+# specify how to quantize the model
+quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype="torch.float16",
+)
+
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1", quantization_config=True, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1")
+
 # Load the tokenizer and model
-model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
-model = AutoModelForCausalLM.from_pretrained(model_id, token=access_token, load_in_4bit=True)
+#model_id = "mistralai/Mixtral-8x7B-v0.1"
+#tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
+#model = AutoModelForCausalLM.from_pretrained(model_id, token=access_token, load_in_4bit=True)
 #model = AutoModelForCausalLM.from_pretrained(model_id, token=access_token)
 # Initialize the quantizer
 #quantizer = bnb.GemmQuantizer(act_bits=8, weight_bits=8)