Update app.py

app.py

@@ -7,11 +7,11 @@ from peft import PeftModel, PeftConfig
 
 
 tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/BhashiniLLM")
-# quantization_config = BitsAndBytesConfig(
-#         load_in_4bit=True,
-#         bnb_4bit_use_double_quant=True,
-#         bnb_4bit_quant_type="nf4",  
-#         bnb_4bit_compute_dtype=torch.float16)
+quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",  
+        bnb_4bit_compute_dtype=torch.float16)
 # model = AutoModelForCausalLM.from_pretrained("FlawedLLM/BhashiniLLM",
 #                                              device_map="auto",
 #                                              quantization_config=quantization_config, 
@@ -35,7 +35,7 @@ tokenizer = AutoTokenizer.from_pretrained("FlawedLLM/BhashiniLLM")
 
 
 config = PeftConfig.from_pretrained("FlawedLLM/BhashiniLLM")
-base_model = AutoModelForCausalLM.from_pretrained("unsloth/llama-3-8b-bnb-4bit",  device_map='auto')
+base_model = AutoModelForCausalLM.from_pretrained("unsloth/llama-3-8b-bnb-4bit",  device_map='auto', quantization_config=quantization_config)
 model = PeftModel.from_pretrained(base_model, "FlawedLLM/BhashiniLLM")