Commit
·
a005314
1
Parent(s):
7b92b5b
Update README.md
Browse files
README.md
CHANGED
|
@@ -80,6 +80,7 @@ pip install -r requirements.txt
|
|
| 80 |
```
|
| 81 |
|
| 82 |
```python
|
|
|
|
| 83 |
from transformers import (
|
| 84 |
AutoTokenizer,
|
| 85 |
AutoModelForCausalLM,
|
|
@@ -88,7 +89,13 @@ tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=T
|
|
| 88 |
tokenizer.padding_side = "left"
|
| 89 |
tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
|
| 90 |
tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
|
| 91 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 92 |
|
| 93 |
HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
|
| 94 |
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
|
|
| 193 |
tokenizer.padding_side = "left"
|
| 194 |
tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
|
| 195 |
tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
|
| 196 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 197 |
|
| 198 |
HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
|
| 199 |
BOT_ROLE_START_TAG = "<|role_start|>bot<|role_end|>"
|
|
|
|
| 80 |
```
|
| 81 |
|
| 82 |
```python
|
| 83 |
+
import torch
|
| 84 |
from transformers import (
|
| 85 |
AutoTokenizer,
|
| 86 |
AutoModelForCausalLM,
|
|
|
|
| 89 |
tokenizer.padding_side = "left"
|
| 90 |
tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
|
| 91 |
tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
|
| 92 |
+
# try 4bit loading if cuda memory not enough
|
| 93 |
+
model = AutoModelForCausalLM.from_pretrained(mode_name_or_path,
|
| 94 |
+
trust_remote_code=True,
|
| 95 |
+
load_in_4bit=False,
|
| 96 |
+
device_map="auto",
|
| 97 |
+
torch_dtype=torch.bfloat16)
|
| 98 |
+
model.eval()
|
| 99 |
|
| 100 |
HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
|
| 101 |
BOT_ROLE_START_TAG = "<|role_start|>bot<|role_end|>"
|
|
|
|
| 200 |
tokenizer.padding_side = "left"
|
| 201 |
tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("<unk>")
|
| 202 |
tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids("</s>")
|
| 203 |
+
# 如果显存不够，可以考虑量化加载
|
| 204 |
+
model = AutoModelForCausalLM.from_pretrained(mode_name_or_path,
|
| 205 |
+
trust_remote_code=True,
|
| 206 |
+
load_in_4bit=False,
|
| 207 |
+
device_map="auto",
|
| 208 |
+
torch_dtype=torch.bfloat16)
|
| 209 |
+
model.eval()
|
| 210 |
|
| 211 |
HUMAN_ROLE_START_TAG = "<|role_start|>human<|role_end|>"
|
| 212 |
BOT_ROLE_START_TAG = "<|role_start|>bot<|role_end|>"
|