Spaces:
Running
Running
Upload chinese_bert.py
Browse files- text/chinese_bert.py +1 -10
text/chinese_bert.py
CHANGED
@@ -1,16 +1,7 @@
|
|
1 |
import torch
|
2 |
-
import sys
|
3 |
from transformers import AutoTokenizer, AutoModelForMaskedLM
|
4 |
|
5 |
-
device = torch.device(
|
6 |
-
"cuda"
|
7 |
-
if torch.cuda.is_available()
|
8 |
-
else (
|
9 |
-
"mps"
|
10 |
-
if sys.platform == "darwin" and torch.backends.mps.is_available()
|
11 |
-
else "cpu"
|
12 |
-
)
|
13 |
-
)
|
14 |
|
15 |
tokenizer = AutoTokenizer.from_pretrained("./bert/chinese-roberta-wwm-ext-large")
|
16 |
model = AutoModelForMaskedLM.from_pretrained("./bert/chinese-roberta-wwm-ext-large").to(device)
|
|
|
1 |
import torch
|
|
|
2 |
from transformers import AutoTokenizer, AutoModelForMaskedLM
|
3 |
|
4 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
5 |
|
6 |
tokenizer = AutoTokenizer.from_pretrained("./bert/chinese-roberta-wwm-ext-large")
|
7 |
model = AutoModelForMaskedLM.from_pretrained("./bert/chinese-roberta-wwm-ext-large").to(device)
|