Update utils.py

utils.py

@@ -18,6 +18,13 @@ model_id = 'meta-llama/Llama-2-7b-chat-hf'
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+bnb_config = transformers.BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type='nf4',
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=bfloat16
+)
+
 model_config = transformers.AutoConfig.from_pretrained(
     model_id,
     use_auth_token=hf_auth
@@ -27,6 +34,7 @@ llm_model = transformers.AutoModelForCausalLM.from_pretrained(
     model_id,
     trust_remote_code=True,
     config=model_config,
+    quantization_config=bnb_config,
     device_map='auto',
     use_auth_token=hf_auth
 )