Update README.md

README.md

@@ -31,12 +31,22 @@ The additional details of the Aquila model will be presented in the official tec
 Aquila2-7B is a base model that can be used for continuation.
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import BitsAndBytesConfig
+
 device = torch.device("cuda")
 model_info = "BAAI/Aquila2-7B"
 tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True)
+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.float16,
+                                                # quantization_config=quantization_config, # Uncomment this line for 4bit quantization
+                                                )
 model.eval()
 model.to(device)
 text = "杭州亚运会的亮点和期待 2023年9月23日至10月8日，杭州将举办第19届亚洲运动会"