|
|
@ -47,11 +47,11 @@ if os.path.exists(txt_path) == False:
|
|
|
|
bert_dir = "%s/3-bert" % (opt_dir)
|
|
|
|
bert_dir = "%s/3-bert" % (opt_dir)
|
|
|
|
os.makedirs(opt_dir, exist_ok=True)
|
|
|
|
os.makedirs(opt_dir, exist_ok=True)
|
|
|
|
os.makedirs(bert_dir, exist_ok=True)
|
|
|
|
os.makedirs(bert_dir, exist_ok=True)
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
if torch.cuda.is_available():
|
|
|
|
device = "cuda:0"
|
|
|
|
device = "cuda:0"
|
|
|
|
elif torch.backends.mps.is_available():
|
|
|
|
elif torch.backends.mps.is_available():
|
|
|
|
device = "mps"
|
|
|
|
device = "mps"
|
|
|
|
else:
|
|
|
|
else:
|
|
|
|
device = "cpu"
|
|
|
|
device = "cpu"
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(bert_pretrained_dir)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(bert_pretrained_dir)
|
|
|
|
bert_model = AutoModelForMaskedLM.from_pretrained(bert_pretrained_dir)
|
|
|
|
bert_model = AutoModelForMaskedLM.from_pretrained(bert_pretrained_dir)
|
|
|
|