Update README.md

README.md

@@ -24,11 +24,21 @@ The additional details of the Aquila model will be presented in the official tec
 
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import BitsAndBytesConfig
 import torch
+
 device = torch.device("cuda:0")
 model_info = "BAAI/AquilaChat2-34B"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+quantization_config=BitsAndBytesConfig(
+                        load_in_4bit=True,
+                        bnb_4bit_use_double_quant=True,
+                        bnb_4bit_quant_type="nf4",
+                        bnb_4bit_compute_dtype=torch.bfloat16,
+                    )
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16
+                                                # quantization_config=quantization_config, # Uncomment this line for 4bit quantization
+                                                )
 model.eval()
 model.to(device)
 text = "请给出10个要到北京旅游的理由。"