File size: 4,028 Bytes
e9057bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import json
import os
import random
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer
class LongBERTConfig(object):
def __init__(self, tokenizer=None):
self.attention_probs_dropout_prob = 0.1
self.hidden_dropout_prob = 0.1
self.hidden_size = 768
self.max_position_embeddings = 70_000
self.num_attention_heads = 12
self.num_hidden_layers = 12
self.pad_token_id = 0
self.vocab_size = len(tokenizer) if tokenizer is not None else None
self.segment_size = [16, 128, 512, 1024, 2048]
self.dilated_rate = [1, 16, 64, 256, 512]
def __call__(self):
return self
def __str__(self):
return str(self.__dict__)
def save_pretrained(self, ckpt="."):
ckpt = os.path.join(ckpt, "config.json")
with open(ckpt, "w") as f:
json.dump(self.__dict__, f)
@classmethod
def from_pretrained(cls, ckpt):
if os.path.isdir(ckpt):
path = os.path.join(ckpt, "config.json")
elif os.path.isfile(ckpt):
path = ckpt
else:
path = hf_hub_download(repo_id=ckpt, filename="config.json")
with open(path, "r") as f:
config_json = json.load(f)
return cls.from_dict(config_json)
@classmethod
def from_dict(cls, _dict):
config = cls()
config.attention_probs_dropout_prob = _dict["attention_probs_dropout_prob"]
config.hidden_dropout_prob = _dict["hidden_dropout_prob"]
config.hidden_size = _dict["hidden_size"]
config.max_position_embeddings = _dict["max_position_embeddings"]
config.num_attention_heads = _dict["num_attention_heads"]
config.num_hidden_layers = _dict["num_hidden_layers"]
config.pad_token_id = _dict["pad_token_id"]
config.vocab_size = _dict["vocab_size"]
config.segment_size = _dict["segment_size"]
config.dilated_rate = _dict["dilated_rate"]
return config
class Config(object):
def __init__(self, args):
# General settings
self.seed = args.seed
self.ver = args.ver
self.use_log = bool(args.use_log)
self.use_tqdm = bool(args.use_tqdm)
self.debug = bool(args.debug)
# Model
backbone = args.backbone
self.tokenizer = AutoTokenizer.from_pretrained(backbone)
self.config = LongBERTConfig(self.tokenizer)
# Data
self.max_len = args.max_len
# Training
self.train_one_part = bool(args.train_one_part)
self.gradient_accumulation_steps = args.gradient_accumulation_steps
self.apex = bool(args.apex)
self.device = torch.device(args.device)
self.nepochs = args.nepochs
self.batch_size = args.batch_size
self.num_workers = os.cpu_count()
# Optimizer
self.lr = args.lr
self.weight_decay = args.weight_decay
self.encoder_lr = self.lr
self.decoder_lr = 1e-3
self.min_lr = args.min_lr
self.eps = 1e-6
self.betas = (0.9, 0.999)
# Scheduler
self.scheduler_type = args.scheduler_type
if self.scheduler_type == "cosine":
self.num_cycles = 0.5
self.num_warmup_steps = args.num_warmup_steps
# Paths
self.train_data_dir = args.train_data_dir
self.valid_data_dir = args.valid_data_dir
self.test_data_dir = args.test_data_dir
self.output_dir = f"model/{self.ver[:-1]}/{self.ver[-1]}"
os.makedirs(self.output_dir, exist_ok=True)
def __str__(self):
return str(self.__dict__)
def set_random_seed(seed):
np.random.seed(seed)
torch.manual_seed(seed)
random.seed(seed)
os.environ["PYTHONHASHSEED"] = str(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
|