Update app.py

app.py

@@ -5,15 +5,24 @@ import os
 
 import torch
 from peft import PeftConfig, PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
 
 base_model_name = "google/gemma-7b"
 adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l5e5-e3-d25-r8"
 
-model = AutoModelForCausalLM.from_pretrained(base_model_name, token=os.environ['HF_TOKEN'], device_map="auto")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    #bnb_4bit_quant_type="nf4",
+    #bnb_4bit_compute_dtype=torch.bfloat16,
+    #bnb_4bit_use_double_quant=True
+)
+
+model = AutoModelForCausalLM.from_pretrained(base_model_name,
+                                             token=os.environ['HF_TOKEN'],
+                                             quantization_config=bnb_config,
+                                             device_map="auto")
 model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
 model.merge_and_unload()