Update README.md

README.md

@@ -81,4 +81,43 @@ with torch.inference_mode():
     outputs = model.generate(**inputs, max_new_tokens=512) 
     result_text = tokenizer.batch_decode(outputs)[0]
     print(result_text)
-```
+```
+
+# Load model in 4bits
+```python
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, BitsAndBytesConfig
+
+# quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
+bnb_config = BitsAndBytesConfig(
+    llm_int8_enable_fp32_cpu_offload=True,
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+# control model memory allocation between devices for low GPU resource (0,cpu)
+device_map = {
+    "transformer.word_embeddings": 0,
+    "transformer.word_embeddings_layernorm": 0,
+    "lm_head": 0,
+    "transformer.h": 0,
+    "transformer.ln_f": 0,
+    "model.embed_tokens": 0,
+    "model.layers":0,
+    "model.norm":0    
+}
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# model use for inference
+model_id="mychen76/mistral7b_ocr_to_json_v1"
+model = AutoModelForCausalLM.from_pretrained(
+    model_id, 
+    trust_remote_code=True,  
+    torch_dtype=torch.float16,
+    quantization_config=bnb_config,
+    device_map=device_map)
+# tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+```