Update app.py

app.py

@@ -31,12 +31,38 @@ model_config = transformers.AutoConfig.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "right"
+#################################################################
+# bitsandbytes parameters
+#################################################################
+
+# Activate 4-bit precision base model loading
+use_4bit = True
+
+# Compute dtype for 4-bit base models
+bnb_4bit_compute_dtype = "float16"
+
+# Quantization type (fp4 or nf4)
+bnb_4bit_quant_type = "nf4"
+
+# Activate nested quantization for 4-bit base models (double quantization)
+use_nested_quant = False
+#################################################################
+# Set up quantization config
+#################################################################
+compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=use_4bit,
+    bnb_4bit_quant_type=bnb_4bit_quant_type,
+    bnb_4bit_compute_dtype=compute_dtype,
+    bnb_4bit_use_double_quant=use_nested_quant,
+)
 
 #############################################################
 # Load pre-trained config
 #################################################################
 model = AutoModelForCausalLM.from_pretrained(
-   "mistralai/Mistral-7B-Instruct-v0.1"
+   "mistralai/Mistral-7B-Instruct-v0.1",quantization_config=bnb_config,
 )
 # Connect query to FAISS index using a retriever
 retriever = db.as_retriever(