Update README.md

README.md

@@ -80,6 +80,7 @@ pip install -r requirements.txt
 ```
 
 ```python
+import torch
 from transformers import (
     AutoTokenizer, 
     AutoModelForCausalLM,
@@ -88,7 +89,13 @@ tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=T
 tokenizer.padding_side = "left"
 tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
 tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
-model = AutoModelForCausalLM.from_pretrained(mode_name_or_path, trust_remote_code=True)
+# try 4bit loading if cuda memory not enough
+model = AutoModelForCausalLM.from_pretrained(mode_name_or_path,
+                                             trust_remote_code=True,
+                                             load_in_4bit=False,
+                                             device_map="auto",
+                                             torch_dtype=torch.bfloat16)
+model.eval()
 
 HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
 BOT_ROLE_START_TAG = "<|role_start|>bot<|role_end|>"
@@ -193,7 +200,13 @@ tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=T
 tokenizer.padding_side = "left"
 tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
 tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
-model = AutoModelForCausalLM.from_pretrained(mode_name_or_path, trust_remote_code=True)
+# 如果显存不够，可以考虑量化加载
+model = AutoModelForCausalLM.from_pretrained(mode_name_or_path,
+                                             trust_remote_code=True,
+                                             load_in_4bit=False,
+                                             device_map="auto",
+                                             torch_dtype=torch.bfloat16)
+model.eval()
 
 HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
 BOT_ROLE_START_TAG = "<|role_start|>bot<|role_end|>"