def bertForTokenClassification(*args, **kwargs):
"""
BertForTokenClassification is a fine-tuning model that includes BertModel
and a token-level classifier on top of the BertModel. Note that the classification
head is only initialized and has to be trained.
The token-level classifier is a linear layer that takes as input the last
hidden state of the sequence.
Args:
num_labels: the number (>=2) of classes for the classifier.
Example:
# Load the tokenizer
>>> import torch
>>> tokenizer = torch.hub.load('huggingface/pytorch-pretrained-BERT', 'bertTokenizer', 'bert-base-cased', do_basic_tokenize=False)
# Prepare tokenized input
>>> text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
>>> tokenized_text = tokenizer.tokenize(text)
>>> indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
>>> segments_ids = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
>>> tokens_tensor = torch.tensor([indexed_tokens])
>>> segments_tensors = torch.tensor([segments_ids])
# Load bertForTokenClassification
>>> model = torch.hub.load('huggingface/pytorch-pretrained-BERT', 'bertForTokenClassification', 'bert-base-cased', num_labels=2)
>>> model.eval()
# Predict the token classification logits
>>> with torch.no_grad():
classif_logits = model(tokens_tensor, segments_tensors)
# Or get the token classification loss
>>> labels = torch.tensor([[0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0]])
>>> classif_loss = model(tokens_tensor, segments_tensors, labels=labels) # set model.train() before if training this loss
"""
model = BertForTokenClassification.from_pretrained(*args, **kwargs)
return model
def __init__(self, language=Language.ENGLISH, num_labels=2, cache_dir="."):
"""
Initializes the classifier and the underlying pre-trained model.
Args:
language (Language, optional): The pre-trained model's language.
The value of this argument determines which BERT model is
used:
Language.ENGLISH: "bert-base-uncased"
Language.ENGLISHCASED: "bert-base-cased"
Language.ENGLISHLARGE: "bert-large-uncased"
Language.ENGLISHLARGECASED: "bert-large-cased"
Language.CHINESE: "bert-base-chinese"
Language.MULTILINGUAL: "bert-base-multilingual-cased"
Defaults to Language.ENGLISH.
num_labels (int, optional): The number of unique labels in the
data. Defaults to 2.
cache_dir (str, optional): Location of BERT's cache directory.
Defaults to ".".
"""
if num_labels < 2:
raise ValueError("Number of labels should be at least 2.")
self.language = language
self.num_labels = num_labels
self.cache_dir = cache_dir
self.model = BertForTokenClassification.from_pretrained(
language, cache_dir=cache_dir, num_labels=num_labels)
self.has_cuda = self.cuda
def __init__(self):
super(Bert_CRF, self).__init__()
self.bert = BertForTokenClassification.from_pretrained(
args.bert_model,
cache_dir=PYTORCH_PRETRAINED_BERT_CACHE,
num_labels=len(args.labels))
self.crf = CRF(len(args.labels))
def bertForTokenClassification(*args, **kwargs):
"""
BertForTokenClassification is a fine-tuning model that includes BertModel
and a token-level classifier on top of the BertModel.
The token-level classifier is a linear layer that takes as input the last
hidden state of the sequence.
"""
model = BertForTokenClassification.from_pretrained(*args, **kwargs)
return model
def bertForTokenClassification(*args, **kwargs):
"""
BertForTokenClassification is a fine-tuning model that includes BertModel
and a token-level classifier on top of the BertModel.
The token-level classifier is a linear layer that takes as input the last
hidden state of the sequence.
Args:
num_labels: the number (>=2) of classes for the classifier.
Example:
>>> torch.hub.load('huggingface/pytorch-pretrained-BERT', 'bertForTokenClassification', 'bert-base-cased', num_labels=2, force_reload=True)
"""
model = BertForTokenClassification.from_pretrained(*args, **kwargs)
return model
def __init__(self, data_dir, bert_model_dir, fine_tuning_model_dir):
self.max_seq_length = 128
task_name = "MSRANER"
eval_batch_size = 32
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = DataProcessor(os.path.join(data_dir, task_name), do_lower_case=True)
processor.get_train_examples()
self.label_list = processor.all_labels
num_labels = len(self.label_list)
self.tokenizer = BertTokenizer.from_pretrained(bert_model_dir, do_lower_case=True)
output_model_file = os.path.join(fine_tuning_model_dir, task_name, "pytorch_model.bin")
# Load a trained model that you have fine-tuned
model_state_dict = torch.load(output_model_file)
self.model = BertForTokenClassification.from_pretrained(bert_model_dir, state_dict=model_state_dict,
num_labels=num_labels)
self.model.to(self.device)
self.model.eval()
self.all_labels = processor.all_labels
def main():
parser = train_opts()
args, _ = parser.parse_known_args()
label_list = [
"O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC",
"I-LOC", "X", "[CLS]", "[SEP]"
]
num_labels = len(label_list) + 1
# Load features
train_features = pd.read_parquet(os.path.join(args.train_feature_dir,
"feature.parquet"),
engine='pyarrow')
input_ids_list = train_features['input_ids'].tolist()
input_mask_list = train_features['input_mask'].tolist()
segment_ids_list = train_features['segment_ids'].tolist()
label_ids_list = train_features['label_ids'].tolist()
all_input_ids = torch.tensor(input_ids_list, dtype=torch.long)
all_input_mask = torch.tensor(input_mask_list, dtype=torch.long)
all_segment_ids = torch.tensor(segment_ids_list, dtype=torch.long)
all_label_ids = torch.tensor(label_ids_list, dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids,
all_label_ids)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not os.path.exists(args.output_model_dir):
os.makedirs(args.output_model_dir)
num_train_optimization_steps = int(
len(train_features) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_features))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(
global_step / num_train_optimization_steps,
args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_model_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_model_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
label_map = {i: label for i, label in enumerate(label_list, 1)}
model_config = {
"bert_model": args.bert_model,
"do_lower": args.do_lower_case,
"max_seq_length": args.max_seq_length,
"num_labels": len(label_list) + 1,
"label_map": label_map
}
json.dump(
model_config,
open(os.path.join(args.output_model_dir, "model_config.json"), "w"))
# Dump data_type.json as a work around until SMT deploys
dct = {
"Id": "ILearnerDotNet",
"Name": "ILearner .NET file",
"ShortName": "Model",
"Description": "A .NET serialized ILearner",
"IsDirectory": False,
"Owner": "Microsoft Corporation",
"FileExtension": "ilearner",
"ContentType": "application/octet-stream",
"AllowUpload": False,
"AllowPromotion": False,
"AllowModelPromotion": True,
"AuxiliaryFileExtension": None,
"AuxiliaryContentType": None
}
with open(os.path.join(args.output_model_dir, 'data_type.json'), 'w') as f:
json.dump(dct, f)
# Dump data.ilearner as a work around until data type design
visualization = os.path.join(args.output_model_dir, "data.ilearner")
with open(visualization, 'w') as file:
file.writelines('{}')
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument("--data_dir",
default=None,
type=str,
required=True,
help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
parser.add_argument("--src_file", default=None, type=str,
help="The input data file name.")
parser.add_argument("--tgt_file", default=None, type=str,
help="The output data file name.")
parser.add_argument("--bert_model", default=None, type=str, required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese.")
parser.add_argument("--config_path", default=None, type=str,
help="Bert config file path.")
parser.add_argument("--output_dir",
default=None,
type=str,
required=True,
help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument("--log_dir",
default='',
type=str,
required=True,
help="The output directory where the log will be written.")
parser.add_argument("--model_recover_path",
default=None,
type=str,
required=True,
help="The file of fine-tuned pretraining model.")
parser.add_argument("--optim_recover_path",
default=None,
type=str,
help="The file of pretraining optimizer.")
# Other parameters
parser.add_argument("--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=64,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate", default=5e-5, type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--label_smoothing", default=0, type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay",
default=0.01,
type=float,
help="The weight decay rate for Adam.")
parser.add_argument("--finetune_decay",
action='store_true',
help="Weight decay to the original weights.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--hidden_dropout_prob", default=0.1, type=float,
help="Dropout rate for hidden states.")
parser.add_argument("--attention_probs_dropout_prob", default=0.1, type=float,
help="Dropout rate for attention probabilities.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument('--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument('--fp16', action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument('--fp32_embedding', action='store_true',
help="Whether to use 32-bit float precision instead of 16-bit for embeddings")
parser.add_argument('--loss_scale', type=float, default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument('--amp', action='store_true',
help="Whether to use amp for fp16")
parser.add_argument('--from_scratch', action='store_true',
help="Initialize parameters with random values (i.e., training from scratch).")
parser.add_argument('--new_segment_ids', action='store_true',
help="Use new segment ids for bi-uni-directional LM.")
parser.add_argument('--new_pos_ids', action='store_true',
help="Use new position ids for LMs.")
parser.add_argument('--tokenized_input', action='store_true',
help="Whether the input is tokenized.")
parser.add_argument('--max_len_a', type=int, default=0,
help="Truncate_config: maximum length of segment A.")
parser.add_argument('--max_len_b', type=int, default=0,
help="Truncate_config: maximum length of segment B.")
parser.add_argument('--trunc_seg', default='',
help="Truncate_config: first truncate segment A/B (option: a, b).")
parser.add_argument('--always_truncate_tail', action='store_true',
help="Truncate_config: Whether we should always truncate tail.")
parser.add_argument("--mask_prob", default=0.15, type=float,
help="Number of prediction is sometimes less than max_pred when sequence is short.")
parser.add_argument("--mask_prob_eos", default=0, type=float,
help="Number of prediction is sometimes less than max_pred when sequence is short.")
parser.add_argument('--max_pred', type=int, default=20,
help="Max tokens of prediction.")
parser.add_argument("--num_workers", default=0, type=int,
help="Number of workers for the data loader.")
parser.add_argument('--mask_source_words', action='store_true',
help="Whether to mask source words for training")
parser.add_argument('--skipgram_prb', type=float, default=0.0,
help='prob of ngram mask')
parser.add_argument('--skipgram_size', type=int, default=1,
help='the max size of ngram mask')
parser.add_argument('--mask_whole_word', action='store_true',
help="Whether masking a whole word.")
parser.add_argument('--do_l2r_training', action='store_true',
help="Whether to do left to right training")
parser.add_argument('--has_sentence_oracle', action='store_true',
help="Whether to have sentence level oracle for training. "
"Only useful for summary generation")
parser.add_argument('--max_position_embeddings', type=int, default=None,
help="max position embeddings")
parser.add_argument('--relax_projection', action='store_true',
help="Use different projection layers for tasks.")
parser.add_argument('--ffn_type', default=0, type=int,
help="0: default mlp; 1: W((Wx+b) elem_prod x);")
parser.add_argument('--num_qkv', default=0, type=int,
help="Number of different <Q,K,V>.")
parser.add_argument('--seg_emb', action='store_true',
help="Using segment embedding for self-attention.")
parser.add_argument('--s2s_special_token', action='store_true',
help="New special tokens ([S2S_SEP]/[S2S_CLS]) of S2S.")
parser.add_argument('--s2s_add_segment', action='store_true',
help="Additional segmental for the encoder of S2S.")
parser.add_argument('--s2s_share_segment', action='store_true',
help="Sharing segment embeddings for the encoder of S2S (used with --s2s_add_segment).")
parser.add_argument('--pos_shift', action='store_true',
help="Using position shift for fine-tuning.")
parser.add_argument('--eval_file', type=str, default="")
args = parser.parse_args()
assert Path(args.model_recover_path).exists(
), "--model_recover_path doesn't exist"
args.output_dir = args.output_dir.replace(
'[PT_OUTPUT_DIR]', os.getenv('PT_OUTPUT_DIR', ''))
args.log_dir = args.log_dir.replace(
'[PT_OUTPUT_DIR]', os.getenv('PT_OUTPUT_DIR', ''))
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(args.log_dir, exist_ok=True)
json.dump(args.__dict__, open(os.path.join(
args.output_dir, 'opt.json'), 'w'), sort_keys=True, indent=2)
if args.local_rank == -1 or args.no_cuda:
device = torch.device(
"cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
dist.init_process_group(backend='nccl')
logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
args.gradient_accumulation_steps))
args.train_batch_size = int(
args.train_batch_size / args.gradient_accumulation_steps)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if args.local_rank not in (-1, 0):
# Make sure only the first process in distributed training will download model & vocab
dist.barrier()
tokenizer = BertTokenizer.from_pretrained(
args.bert_model, do_lower_case=args.do_lower_case)
if args.max_position_embeddings:
tokenizer.max_len = args.max_position_embeddings
data_tokenizer = WhitespaceTokenizer() if args.tokenized_input else tokenizer
if args.local_rank == 0:
dist.barrier()
if args.do_train:
print("Loading Train Dataset", args.data_dir)
bi_uni_pipeline = [Preprocess4CoNLL2003(args.max_pred, args.mask_prob, list(tokenizer.vocab.keys(
)), tokenizer.convert_tokens_to_ids, args.max_seq_length, new_segment_ids=args.new_segment_ids, truncate_config={'max_len_a': args.max_len_a, 'max_len_b': args.max_len_b, 'trunc_seg': args.trunc_seg, 'always_truncate_tail': args.always_truncate_tail}, mask_source_words=args.mask_source_words, skipgram_prb=args.skipgram_prb, skipgram_size=args.skipgram_size, mask_whole_word=args.mask_whole_word, mode="s2s", has_oracle=args.has_sentence_oracle, num_qkv=args.num_qkv, s2s_special_token=args.s2s_special_token, s2s_add_segment=args.s2s_add_segment, s2s_share_segment=args.s2s_share_segment, pos_shift=args.pos_shift)]
file_oracle = None
if args.has_sentence_oracle:
file_oracle = os.path.join(args.data_dir, 'train.oracle')
fn_src = os.path.join(
args.data_dir, args.src_file if args.src_file else 'train.src')
fn_tgt = os.path.join(
args.data_dir, args.tgt_file if args.tgt_file else 'train.tgt')
train_dataset = CoNLL2003Dataset(
fn_src, fn_tgt, args.train_batch_size, data_tokenizer, args.max_seq_length, file_oracle=file_oracle, bi_uni_pipeline=bi_uni_pipeline)
if args.local_rank == -1:
train_sampler = RandomSampler(train_dataset, replacement=False)
_batch_size = args.train_batch_size
else:
train_sampler = DistributedSampler(train_dataset)
_batch_size = args.train_batch_size // dist.get_world_size()
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=_batch_size, sampler=train_sampler,
num_workers=args.num_workers, collate_fn=seq2seq_loader.batch_list_to_batch_tensors, pin_memory=False)
# note: args.train_batch_size has been changed to (/= args.gradient_accumulation_steps)
# t_total = int(math.ceil(len(train_dataset.ex_list) / args.train_batch_size)
t_total = int(len(train_dataloader) * args.num_train_epochs /
args.gradient_accumulation_steps)
amp_handle = None
if args.fp16 and args.amp:
from apex import amp
amp_handle = amp.init(enable_caching=True)
logger.info("enable fp16 with amp")
# Prepare model
recover_step = _get_max_epoch_model(args.output_dir)
cls_num_labels = 2
type_vocab_size = 6 + \
(1 if args.s2s_add_segment else 0) if args.new_segment_ids else 2
num_sentlvl_labels = 2 if args.has_sentence_oracle else 0
relax_projection = 4 if args.relax_projection else 0
if args.local_rank not in (-1, 0):
# Make sure only the first process in distributed training will download model & vocab
dist.barrier()
if (recover_step is None) and (args.model_recover_path is None):
# if _state_dict == {}, the parameters are randomly initialized
# if _state_dict == None, the parameters are initialized with bert-init
_state_dict = {} if args.from_scratch else None
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=10)
global_step = 0
else:
if recover_step:
logger.info("***** Recover model: %d *****", recover_step)
model_recover = torch.load(os.path.join(
args.output_dir, "model.{0}.bin".format(recover_step)), map_location='cpu')
# recover_step == number of epochs
global_step = math.floor(
recover_step * t_total / args.num_train_epochs)
elif args.model_recover_path:
logger.info("***** Recover model: %s *****",
args.model_recover_path)
model_recover = torch.load(
args.model_recover_path, map_location='cpu')
global_step = 0
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=10)
if args.local_rank == 0:
dist.barrier()
if args.fp16:
model.half()
if args.fp32_embedding:
model.bert.embeddings.word_embeddings.float()
model.bert.embeddings.position_embeddings.float()
model.bert.embeddings.token_type_embeddings.float()
model.to(device)
if args.local_rank != -1:
try:
from torch.nn.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError("DistributedDataParallel")
model = DDP(model, device_ids=[
args.local_rank], output_device=args.local_rank, find_unused_parameters=True)
elif n_gpu > 1:
# model = torch.nn.DataParallel(model)
model = DataParallelImbalance(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(
nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(
nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
if args.fp16:
try:
# from apex.optimizers import FP16_Optimizer
from pytorch_pretrained_bert.optimization_fp16 import FP16_Optimizer_State
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer_State(
optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer_State(
optimizer, static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=t_total)
if recover_step:
logger.info("***** Recover optimizer: %d *****", recover_step)
optim_recover = torch.load(os.path.join(
args.output_dir, "optim.{0}.bin".format(recover_step)), map_location='cpu')
if hasattr(optim_recover, 'state_dict'):
optim_recover = optim_recover.state_dict()
optimizer.load_state_dict(optim_recover)
if args.loss_scale == 0:
logger.info("***** Recover optimizer: dynamic_loss_scale *****")
optimizer.dynamic_loss_scale = True
logger.info("***** CUDA.empty_cache() *****")
torch.cuda.empty_cache()
if args.do_train:
logger.info("***** Running training *****")
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", t_total)
model.train()
if recover_step:
start_epoch = recover_step+1
else:
start_epoch = 1
for i_epoch in trange(start_epoch, int(args.num_train_epochs)+1, desc="Epoch", disable=args.local_rank not in (-1, 0)):
if args.local_rank != -1:
train_sampler.set_epoch(i_epoch)
iter_bar = tqdm(train_dataloader, desc='Iter (loss=X.XXX)',
disable=args.local_rank not in (-1, 0))
for step, batch in enumerate(iter_bar):
batch = [
t.to(device) if t is not None else None for t in batch]
if args.has_sentence_oracle:
input_ids, segment_ids, input_mask, mask_qkv, lm_label_ids, masked_pos, masked_weights, is_next, task_idx, oracle_pos, oracle_weights, oracle_labels = batch
else:
input_ids, segment_ids, input_mask, mask_qkv, lm_label_ids, masked_pos, masked_weights, is_next, task_idx, label_ids, valid_length = batch
oracle_pos, oracle_weights, oracle_labels = None, None, None
print("labels_ids_num:")
print(Counter( label_ids.cpu().detach().numpy()))
loss = model(input_ids, segment_ids, input_mask, label_ids, mask_qkv, task_idx)
if n_gpu > 1: # mean() to average on multi-gpu.
loss = loss.mean()
# logging for each step (i.e., before normalization by args.gradient_accumulation_steps)
iter_bar.set_description('Iter (loss=%5.3f)' % loss.item())
# ensure that accumlated gradients are normalized
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
if amp_handle:
amp_handle._clear_cache()
else:
loss.backward()
if (step + 1) % args.gradient_accumulation_steps == 0:
lr_this_step = args.learning_rate * \
warmup_linear(global_step/t_total,
args.warmup_proportion)
if args.fp16:
# modify learning rate with special warm up BERT uses
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model
if (args.local_rank == -1 or torch.distributed.get_rank() == 0):
logger.info(
"** ** * Saving fine-tuned model and optimizer ** ** * ")
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(
args.output_dir, "model.{0}.bin".format(i_epoch))
torch.save(model_to_save.state_dict(), output_model_file)
output_optim_file = os.path.join(
args.output_dir, "optim.{0}.bin".format(i_epoch))
torch.save(optimizer.state_dict(), output_optim_file)
logger.info("***** CUDA.empty_cache() *****")
torch.cuda.empty_cache()
if args.do_eval:
print("Loading Eval Dataset", args.data_dir)
bi_uni_pipeline = [Preprocess4CoNLL2003(args.max_pred, args.mask_prob, list(tokenizer.vocab.keys(
)), tokenizer.convert_tokens_to_ids, args.max_seq_length, new_segment_ids=args.new_segment_ids, truncate_config={'max_len_a': args.max_len_a, 'max_len_b': args.max_len_b, 'trunc_seg': args.trunc_seg, 'always_truncate_tail': args.always_truncate_tail}, mask_source_words=args.mask_source_words, skipgram_prb=args.skipgram_prb, skipgram_size=args.skipgram_size, mask_whole_word=args.mask_whole_word, mode="s2s", has_oracle=args.has_sentence_oracle, num_qkv=args.num_qkv, s2s_special_token=args.s2s_special_token, s2s_add_segment=args.s2s_add_segment, s2s_share_segment=args.s2s_share_segment, pos_shift=args.pos_shift)]
file_oracle = None
fn_src = os.path.join(args.data_dir, args.eval_file)
eval_dataset = CoNLL2003Dataset(
fn_src, fn_tgt, args.train_batch_size, data_tokenizer, args.max_seq_length, file_oracle=file_oracle, bi_uni_pipeline=bi_uni_pipeline)
eval_sampler = SequentialSampler(eval_dataset, replacement=False)
_batch_size = args.train_batch_size
eval_dataloader = torch.utils.data.DataLoader(eval_dataset, batch_size=_batch_size, sampler=eval_sampler,
num_workers=args.num_workers, collate_fn=seq2seq_loader.batch_list_to_batch_tensors, pin_memory=False)
# Prepare model
recover_step = _get_max_epoch_model(args.output_dir)
cls_num_labels = 2
type_vocab_size = 6 + \
(1 if args.s2s_add_segment else 0) if args.new_segment_ids else 2
num_sentlvl_labels = 2 if args.has_sentence_oracle else 0
relax_projection = 4 if args.relax_projection else 0
if args.local_rank not in (-1, 0):
# Make sure only the first process in distributed training will download model & vocab
dist.barrier()
if (recover_step is None) and (args.model_recover_path is None):
# if _state_dict == {}, the parameters are randomly initialized
# if _state_dict == None, the parameters are initialized with bert-init
_state_dict = {} if args.from_scratch else None
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=10)
global_step = 0
else:
if recover_step:
logger.info("***** Recover model: %d *****", recover_step)
model_recover = torch.load(os.path.join(
args.output_dir, "model.{0}.bin".format(recover_step)), map_location='cpu')
# recover_step == number of epochs
global_step = math.floor(
recover_step * t_total / args.num_train_epochs)
elif args.model_recover_path:
logger.info("***** Recover model: %s *****",
args.model_recover_path)
model_recover = torch.load(
args.model_recover_path, map_location='cpu')
global_step = 0
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=10)
if args.local_rank == 0:
dist.barrier()
if args.fp16:
model.half()
if args.fp32_embedding:
model.bert.embeddings.word_embeddings.float()
model.bert.embeddings.position_embeddings.float()
model.bert.embeddings.token_type_embeddings.float()
model.to(device)
if args.local_rank != -1:
try:
from torch.nn.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError("DistributedDataParallel")
model = DDP(model, device_ids=[
args.local_rank], output_device=args.local_rank, find_unused_parameters=True)
elif n_gpu > 1:
# model = torch.nn.DataParallel(model)
model = DataParallelImbalance(model)
logger.info("***** CUDA.empty_cache() *****")
torch.cuda.empty_cache()
logger.info("***** Running Evaluation *****")
logger.info(" Batch size = %d", args.train_batch_size)
model.eval()
iter_bar = tqdm(eval_dataloader, desc="Evaluating")
acc_score = 0.0000
total_t = 0
for step, batch in enumerate(iter_bar):
batch = [
t.to(device) if t is not None else None for t in batch]
input_ids, segment_ids, input_mask, mask_qkv, lm_label_ids, masked_pos, masked_weights, is_next, task_idx, label_ids, valid_length = batch
oracle_pos, oracle_weights, oracle_labels = None, None, None
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask, mask_qkv=mask_qkv, task_idx=task_idx)
valid_length = valid_length.view(-1)
logits = logits.view(-1, input_ids.size(1), 9)
for i in range(input_ids.size(0)):
valid_len = valid_length[i]
logits_i = logits[i]
pred_i = torch.argmax(logits_i, dim=-1)[:valid_len]
labels_i = label_ids[i][:valid_len]
acc_t = acc(pred_i, labels_i)
acc_score += acc_t
total_t += 1
print("acc score:", acc_score/len(total_t))
# all_test_len = [len(item.label) for item in test_examples]
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese."
)
parser.add_argument("--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train.")
parser.add_argument(
"--output_file",
default=None,
type=str,
required=True,
help="The output file where the model checkpoints will be written.")
## Other parameters
parser.add_argument(
"--max_seq_length",
default=512,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument(
"--bert_tokenizer",
default="bert-base-uncased",
type=str,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-base-multilingual, bert-base-chinese."
)
parser.add_argument("--do_train",
default=False,
action='store_true',
help="Whether to run training.")
parser.add_argument("--train_file",
default="",
type=str,
help="The path of train file.")
parser.add_argument("--do_eval",
default=False,
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--dev_file",
default="",
type=str,
help="The path of eval file.")
parser.add_argument("--do_test",
default=False,
action='store_true',
help="Whether to run eval on the test set.")
parser.add_argument("--test_file",
default="",
type=str,
help="The path of test file.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=32,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
default=False,
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumualte before performing a backward/update pass."
)
parser.add_argument(
'--optimize_on_cpu',
default=False,
action='store_true',
help=
"Whether to perform optimization and keep the optimizer averages on CPU"
)
parser.add_argument(
'--fp16',
default=False,
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=128,
help=
'Loss scaling, positive power of 2 values can improve fp16 convergence.'
)
parser.add_argument("--model_load",
default="",
type=str,
help="The path of model state.")
parser.add_argument("--val_each",
default=500000,
type=int,
help="every nth iteration to do eval")
args = parser.parse_args()
processors = {
"emw": EmwProcessor,
}
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
if args.fp16:
logger.info(
"16-bits training currently not supported in distributed training"
)
args.fp16 = False # (see https://github.com/pytorch/pytorch/pull/13496)
logger.info("device %s n_gpu %d distributed training %r", device, n_gpu,
bool(args.local_rank != -1))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = int(args.train_batch_size /
args.gradient_accumulation_steps)
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
# if not args.do_train and not args.do_eval and not args.do_test:
# raise ValueError("At least one of `do_train` or `do_eval` or `do_test` must be True.")
#if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
# raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
#os.makedirs(args.output_dir, exist_ok=True)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer)
label_list = processor.get_labels()
train_examples = None
num_train_steps = None
if args.do_train:
train_examples = processor.get_examples(args.train_file)
random.shuffle(train_examples)
num_train_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps * args.num_train_epochs)
train_dataloader = DataLoader(dataset=HyperpartisanData(
train_examples, label_list, args.max_seq_length, tokenizer),
batch_size=args.train_batch_size)
if args.do_eval:
eval_examples = processor.get_examples(args.dev_file)
eval_dataloader = DataLoader(dataset=HyperpartisanData(
eval_examples, label_list, args.max_seq_length, tokenizer),
batch_size=args.train_batch_size)
if args.do_test:
test_examples = processor.get_examples(args.test_file)
test_dataloader = DataLoader(dataset=HyperpartisanData(
test_examples, label_list, args.max_seq_length, tokenizer),
batch_size=args.train_batch_size)
# Prepare model
num_labels = len(label_list)
model = BertForTokenClassification.from_pretrained(
args.bert_model, PYTORCH_PRETRAINED_BERT_CACHE, num_labels=num_labels)
if args.model_load != "":
model.load_state_dict(torch.load(args.model_load))
logger.info("Model state has been loaded.")
model.to(device)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
if args.fp16:
param_optimizer = [(n, param.clone().detach().to('cpu').float().requires_grad_()) \
for n, param in model.named_parameters()]
elif args.optimize_on_cpu:
param_optimizer = [(n, param.clone().detach().to('cpu').requires_grad_()) \
for n, param in model.named_parameters()]
else:
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'gamma', 'beta']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay_rate':
0.0
}]
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_steps)
idtolabel = {}
for (i, label) in enumerate(label_list):
idtolabel[i] = label
global_step = 0
# best_mcc = 0.0
best_f1 = 0.0
if args.do_train:
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_steps)
model.train()
for epoch_num in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss, _ = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.fp16 and args.loss_scale != 1.0:
# rescale loss for fp16 training
# see https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html
loss = loss * args.loss_scale
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16 or args.optimize_on_cpu:
if args.fp16 and args.loss_scale != 1.0:
# scale down gradients for fp16 training
for param in model.parameters():
if param.grad is not None:
param.grad.data = param.grad.data / args.loss_scale
is_nan = set_optimizer_params_grad(
param_optimizer,
model.named_parameters(),
test_nan=True)
if is_nan:
logger.info(
"FP16 TRAINING: Nan in gradients, reducing loss scaling"
)
args.loss_scale = args.loss_scale / 2
model.zero_grad()
continue
optimizer.step()
copy_optimizer_params_to_model(
model.named_parameters(), param_optimizer)
else:
optimizer.step()
model.zero_grad()
global_step += 1
if global_step % args.val_each == 0:
if args.do_eval:
model.eval()
# total_rates = np.array([0,0,0,0])
all_preds = np.array([])
all_label_ids = np.array([])
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss, logits = model(
input_ids, segment_ids, input_mask,
label_ids)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
# logger.info(logits)
eval_loss += tmp_eval_loss.mean().item()
logits = np.argmax(logits, axis=-1).reshape(-1)
label_ids = label_ids.reshape(-1)
# logger.info(logits)
all_preds = np.append(all_preds,
logits[label_ids != -1])
all_label_ids = np.append(
all_label_ids, label_ids[label_ids != -1])
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = accuracy(all_preds, all_label_ids)
# precision, recall, f1, _ = precision_recall_fscore_support(all_label_ids, all_preds, average="macro", labels=list(range(0,num_labels)))
precision, recall, f1 = evaluate(
[idtolabel[x] for x in all_label_ids.tolist()],
[idtolabel[x] for x in all_preds.tolist()])
mcc = matthews_corrcoef(all_preds, all_label_ids)
result = {
"eval_loss": eval_loss,
"eval_accuracy": eval_accuracy,
"precision": precision,
"recall": recall,
"f1": f1,
"mcc": mcc
}
# balanced_acc, f1_neg, f1_pos, mcc, _, _ = get_scores(total_rates.tolist())
# result = {'eval_loss': eval_loss,
# 'eval_accuracy': eval_accuracy,
# 'global_step': global_step,
# 'balanced_accuracy' : balanced_acc,
# 'f1_neg' : f1_neg,
# 'f1_pos' : f1_pos,
# 'mcc' : mcc,
# 'loss': tr_loss/nb_tr_steps}
# if best_mcc < mcc:
# best_mcc = mcc
if best_f1 < f1:
best_f1 = f1
logger.info("Saving model...")
model_to_save = model.module if hasattr(
model,
'module') else model # To handle multi gpu
torch.save(model_to_save.state_dict(),
args.output_file)
for key in sorted(result.keys()):
logger.info(" %s = %.4f", key, result[key])
model.train() # back to training
if args.do_eval:
model.eval()
# total_rates = np.array([0,0,0,0])
all_preds = np.array([])
all_label_ids = np.array([])
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss, logits = model(input_ids, segment_ids,
input_mask, label_ids)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
logits = np.argmax(logits, axis=-1).reshape(-1)
label_ids = label_ids.reshape(-1)
all_preds = np.append(all_preds, logits[label_ids != -1])
all_label_ids = np.append(all_label_ids,
label_ids[label_ids != -1])
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
eval_accuracy = accuracy(all_preds, all_label_ids)
# precision, recall, f1, _ = precision_recall_fscore_support(all_label_ids, all_preds, average="macro", labels=list(range(0,num_labels)))
precision, recall, f1 = evaluate(
[idtolabel[x] for x in all_label_ids.tolist()],
[idtolabel[x] for x in all_preds.tolist()])
mcc = matthews_corrcoef(all_preds, all_label_ids)
result = {
"eval_loss": eval_loss,
"eval_accuracy": eval_accuracy,
"precision": precision,
"recall": recall,
"f1": f1,
"mcc": mcc
}
# balanced_acc, f1_neg, f1_pos, mcc, _, _ = get_scores(total_rates.tolist())
# result = {'eval_loss': eval_loss,
# 'eval_accuracy': eval_accuracy,
# 'global_step': global_step,
# 'balanced_accuracy' : balanced_acc,
# 'f1_neg' : f1_neg,
# 'f1_pos' : f1_pos,
# 'mcc' : mcc,
# 'loss': tr_loss/nb_tr_steps}
# if best_mcc < mcc:
# best_mcc = mcc
if best_f1 < f1:
best_f1 = f1
logger.info("Saving model...")
model_to_save = model.module if hasattr(
model, 'module') else model # To handle multi gpu
torch.save(model_to_save.state_dict(), args.output_file)
logger.info("***** Epoch " + str(epoch_num + 1) + " *****")
for key in sorted(result.keys()):
logger.info(" %s = %.4f", key, result[key])
model.train() # back to training
# for input_ids, input_mask, segment_ids, label_ids in train_dataloader:
# with open("sst-train.tsv", "a") as f:
# for i in range(len(input_ids)):
# f.write(str(input_ids[i].numpy().tolist()) + "\t" + str(input_mask[i].numpy().tolist()) + "\t" + str(segment_ids[i].numpy().tolist()) + "\t" + str(label_ids[i].numpy().tolist()) + "\n")
# for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
# with open("sst-dev.tsv", "a") as f:
# for i in range(len(input_ids)):
# f.write(str(input_ids[i].numpy().tolist()) + "\t" + str(input_mask[i].numpy().tolist()) + "\t" + str(segment_ids[i].numpy().tolist()) + "\t" + str(label_ids[i].numpy().tolist()) + "\n")
if args.do_test:
model = BertForTokenClassification.from_pretrained(
args.bert_model,
PYTORCH_PRETRAINED_BERT_CACHE,
num_labels=num_labels)
model.load_state_dict(torch.load(args.output_file))
model.to(device)
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(
model,
device_ids=[args.local_rank],
output_device=args.local_rank)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
logger.info("***** Running evaluation on test *****")
logger.info(" Num examples = %d", len(test_examples))
logger.info(" Batch size = %d", args.train_batch_size)
model.eval()
# total_rates = np.array([0,0,0,0])
all_preds = np.array([])
all_label_ids = np.array([])
test_loss, test_accuracy = 0, 0
nb_test_steps, nb_test_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in test_dataloader:
# with open("sst-test.tsv", "a") as f:
# for i in range(len(input_ids)):
# f.write(str(input_ids[i].numpy().tolist()) + "\t" + str(input_mask[i].numpy().tolist()) + "\t" + str(segment_ids[i].numpy().tolist()) + "\t" + str(label_ids[i].numpy().tolist()) + "\n")
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_test_loss, logits = model(input_ids, segment_ids,
input_mask, label_ids)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
test_loss += tmp_test_loss.mean().item()
logits = np.argmax(logits, axis=-1).reshape(-1)
label_ids = label_ids.reshape(-1)
all_preds = np.append(all_preds, logits[label_ids != -1])
all_label_ids = np.append(all_label_ids,
label_ids[label_ids != -1])
nb_test_steps += 1
test_loss = test_loss / nb_test_steps
test_accuracy = accuracy(all_preds, all_label_ids)
# precision, recall, f1, _ = precision_recall_fscore_support(all_label_ids, all_preds, average="macro", labels=list(range(0,num_labels)))
precision, recall, f1 = evaluate(
[idtolabel[x] for x in all_label_ids.tolist()],
[idtolabel[x] for x in all_preds.tolist()],
verbose=True)
mcc = matthews_corrcoef(all_label_ids, all_preds)
result = {
"test_loss": test_loss,
"test_accuracy": test_accuracy,
"precision": precision,
"recall": recall,
"f1": f1,
"mcc": mcc
}
# balanced_acc, f1_neg, f1_pos, mcc, recall_pos, precision_pos = get_scores(total_rates.tolist())
# result = {'test_loss': test_loss,
# 'test_accuracy': test_accuracy,
# 'global_step': global_step,
# 'balanced_accuracy' : balanced_acc,
# 'f1_neg' : f1_neg,
# 'f1_pos' : f1_pos,
# 'recall_pos' : recall_pos,
# 'precision_pos' : precision_pos,
# 'mcc' : mcc}
logger.info("***** Test Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %.4f", key, result[key])
sentdb = data.SentDB(args.sent_fi, args.tag_fi, tokenizer, args.val_sent_fi,
args.val_tag_fi, lower=args.lower,
align_strat=args.align_strat, parampred=True)
sentdb.make_minibatches(args.bsz, None)
sentdb.make_minibatches(args.bsz, None, val=True)
# if args.db_fi is not None:
# print("saving db to", args.db_fi)
# sentdb.save(args.db_fi)
idx2tag = sentdb.tagtypes
if len(args.train_from) > 0:
print("loading model from", args.train_from)
saved_stuff = torch.load(args.train_from)
saved_args = saved_stuff["opt"]
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=len(sentdb.tagtypes),
cache_dir=CACHEDIR)
model.load_state_dict(saved_stuff["state_dict"])
else:
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=len(sentdb.tagtypes),
cache_dir=CACHEDIR)
model = model.to(device)
model.dropout.p = args.drop
if args.just_eval is not None:
import sys
# pos_c2f = {'ADP': 'IN', 'DET': 'DT', 'NOUN': 'NN', 'NUM': 'CD', '.': ',', 'PRT': 'TO',
# 'VERB': 'VBD', 'CONJ': 'CC', 'ADV': 'RB', 'PRON': 'PRP', 'ADJ': 'JJ', 'X': 'FW'}
def main():
parser = argparse.ArgumentParser()
# # Required parameters
parser.add_argument("--data_dir",
default='./msra_ner',
type=str,
help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
parser.add_argument("--bert_model", default='../bert-base-chinese/', type=str,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument("--task_name",
default='ner',
type=str,
help="The name of the task to train.")
parser.add_argument("--output_dir",
default='./out/',
type=str,
help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument('--cuda_device', default=None, type=str, help='e.g. 0')
## Other parameters
parser.add_argument("--cache_dir",
default="",
type=str,
help="Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument("--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
default=True,
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
default=True,
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument('--overwrite_output_dir',
action='store_true',
help="Overwrite the content of the output directory")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument('--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument('--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument('--loss_scale',
type=float, default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument('--server_ip', type=str, default='', help="Can be used for distant debugging.")
parser.add_argument('--server_port', type=str, default='', help="Can be used for distant debugging.")
args = parser.parse_args()
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
ptvsd.wait_for_attach()
if args.cuda_device:
device = torch.device("cuda:%s" % args.cuda_device)
n_gpu = 1
else:
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
args.device = device
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError("At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train and not args.overwrite_output_dir:
raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
os.makedirs(args.output_dir)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]() # nerProcessor
# output_mode = output_modes[task_name]
label_list = processor.get_labels()
num_labels = len(label_list) + 1
if args.local_rank not in [-1, 0]:
torch.distributed.barrier() # Make sure only the first process in distributed training will download model & vocab
stime = time.time()
tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
model = BertForTokenClassification.from_pretrained(args.bert_model, num_labels=num_labels)
etime = time.time()
logger.info("Model loaded in %s seconds" % (etime-stime))
model.to(device)
# gpu_handle = nvmlDeviceGetHandleByIndex(int(args.cuda_device))
# gpu_info = nvmlDeviceGetMemoryInfo(gpu_handle)
# logger.info("GPU usage:%s" % (gpu_info.used / 2 ** 30))
if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
# Prepare data loader
train_examples = processor.get_train_examples(args.data_dir)
cached_train_features_file = os.path.join(args.data_dir, 'train_{0}_{1}_{2}'.format(
list(filter(None, args.bert_model.split('/'))).pop(),
str(args.max_seq_length),
str(task_name)))
try:
with open(cached_train_features_file, "rb") as reader:
train_features = pickle.load(reader)
except:
train_features = convert_examples_to_features_ner(
train_examples, label_list, args.max_seq_length, tokenizer)
# if args.local_rank == -1 or torch.distributed.get_rank() == 0:
# logger.info(" Saving train features into cached file %s", cached_train_features_file)
# with open(cached_train_features_file, "wb") as writer:
# pickle.dump(train_features, writer)
all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)
num_train_optimization_steps = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale)
warmup_linear = WarmupLinearSchedule(warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
model.train()
# gpu_handle = nvmlDeviceGetHandleByIndex(int(args.cuda_device))
# gpu_info = nvmlDeviceGetMemoryInfo(gpu_handle)
# logger.info("Model start training, GPU usage:%s" % (gpu_info.used / 2 ** 30))
for _ in trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0]):
epoch_stime = time.time()
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
# define a new function to compute loss values for both output_modes
loss = model(input_ids=input_ids, token_type_ids=segment_ids, attention_mask=input_mask,
labels=label_ids)
# loss_fct = CrossEntropyLoss()
# loss = loss_fct(logits.view(-1, num_labels), label_ids.view(-1))
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear.get_lr(global_step, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# if args.local_rank in [-1, 0]:
# tb_writer.add_scalar('lr', optimizer.get_lr()[0], global_step)
# tb_writer.add_scalar('loss', loss.item(), global_step)
# epoch_etime = time.time()
# gpu_handle = nvmlDeviceGetHandleByIndex(int(args.cuda_device))
# gpu_info = nvmlDeviceGetMemoryInfo(gpu_handle)
# logger.info(
# "Training epoch cost %s seconds, GPU usage:%s" % (epoch_etime - epoch_stime, gpu_info.used / 2 ** 30))
# logits = model(input_ids=all_input_ids, token_type_ids=all_segment_ids, attention_mask=all_input_mask,
# labels=all_label_ids)
pass
train_dataloader = get_pytorch_dataloader(train_features, batch_size)
dev_dataloader = get_pytorch_dataloader(dev_features, batch_size)
test_dataloader = get_pytorch_dataloader(test_features, batch_size)
#%%
'''
##### Use only BertForTokenClassification #####
'''
print('*** Use only BertForTokenClassification ***')
if True and os.path.exists(output_dir+'/ner_bert_checkpoint.pt'):
checkpoint = torch.load(output_dir+'/ner_bert_checkpoint.pt', map_location='cpu')
start_epoch = checkpoint['epoch']+1
valid_acc_prev = checkpoint['valid_acc']
model = BertForTokenClassification.from_pretrained(
'bert-base-uncased', state_dict=checkpoint['model_state'], num_labels=len(label_list))
else:
start_epoch = 0
valid_acc_prev = 0
model = BertForTokenClassification.from_pretrained(
'bert-base-uncased', num_labels=len(label_list))
model.to(device)
# Prepare optimizer
named_params = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in named_params if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in named_params if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
def main():
parser = get_parser()
args = parser.parse_args()
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
processors = {"ner": NerProcessor}
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels()
num_labels = len(label_list) + 1
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
train_examples = None
num_train_optimization_steps = None
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor(
[f.input_ids for f in train_features],
dtype=torch.long) # words converted to ids by tokenizer
all_input_mask = torch.tensor(
[f.input_mask for f in train_features],
dtype=torch.long) # mask keeps everything
all_segment_ids = torch.tensor(
[f.segment_ids for f in train_features],
dtype=torch.long) # segment ids are all 0
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.long)
# labels are ["O", "B-MISC", "I-MISC", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC", "X", "[CLS]", "[SEP]"]
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train(
) # this just sets the model to be in training mode (is method on pytorch nn.Module)
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(
global_step / num_train_optimization_steps,
args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
label_map = {i: label for i, label in enumerate(label_list, 1)}
model_config = {
"bert_model": args.bert_model,
"do_lower": args.do_lower_case,
"max_seq_length": args.max_seq_length,
"num_labels": len(label_list) + 1,
"label_map": label_map
}
json.dump(
model_config,
open(os.path.join(args.output_dir, "model_config.json"), "w"))
# Load a trained model and config that you have fine-tuned
else:
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
model.to(device)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
eval_examples = processor.get_dev_examples(args.data_dir)
eval_features = convert_examples_to_features(eval_examples, label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
y_true = []
y_pred = []
label_map = {i: label for i, label in enumerate(label_list, 1)}
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask)
logits = torch.argmax(F.log_softmax(logits, dim=2), dim=2)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
input_mask = input_mask.to('cpu').numpy()
for i, mask in enumerate(input_mask):
temp_1 = []
temp_2 = []
for j, m in enumerate(mask):
if j == 0:
continue
if m:
if label_map[label_ids[i][j]] != "X":
temp_1.append(label_map[label_ids[i][j]])
temp_2.append(label_map[logits[i][j]])
else:
temp_1.pop()
temp_2.pop()
break
y_true.append(temp_1)
y_pred.append(temp_2)
report = classification_report(y_true, y_pred, digits=4)
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
logger.info("\n%s", report)
writer.write(report)
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument("--train_file",
default='../../data/eng-2015.conll',
type=str,
required=True,
help="train file path")
parser.add_argument("--dev_file",
default='../../data/eng-2016.conll',
type=str,
required=True,
help="dev file path")
parser.add_argument("--bert_model", default='bert-base-cased', type=str, required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument("--finetune_dir",
default='NER_BERT',
type=str,
required=False,
help="The output")
parser.add_argument("--output_dir",
default='NER_BERT',
type=str,
required=True,
help="The output directory where the model predictions and checkpoints will be written.")
## Other parameters
parser.add_argument("--cache_dir",
default="",
type=str,
help="Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument("--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_finetune",
action='store_true',
help="Whether to run finetuning.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument("--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=32,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument("--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument('--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument('--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument('--loss_scale',
type=float, default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument('--server_ip', type=str, default='', help="Can be used for distant debugging.")
parser.add_argument('--server_port', type=str, default='', help="Can be used for distant debugging.")
args = parser.parse_args()
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info("device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".format(
device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError("At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train:
raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case, do_basic_tokenize=False)
label_list = get_labels()
num_labels = len(label_list)
train_examples = read_ner_example(args.train_file, args.do_lower_case)
num_train_optimization_steps = None
if args.do_train:
#train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size / args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size()
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels = num_labels)
if args.fp16:
model.half()
if args.do_finetune:
if not os.path.exists(args.finetune_dir) and not os.listdir(args.finetune_dir):
raise ValueError("Finetune directory ({}) is empty.".format(args.finetune_dir))
finetune_model_file = os.path.join(args.finetune_dir, WEIGHTS_NAME)
finetune_config_file = os.path.join(args.finetune_dir, CONFIG_NAME)
config = BertConfig(finetune_config_file)
#model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(finetune_model_file))
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
train_features = convert_examples_to_features(
train_examples, label_list, args.max_seq_length, tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.long)
all_label_masks = torch.tensor([f.label_mask for f in train_features], dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids, all_label_masks)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids, label_masks = batch
loss = model(input_ids, segment_ids, input_mask, label_ids, label_masks)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(global_step/num_train_optimization_steps, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
# Load a trained model and config that you have fine-tuned
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
else:
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
#model = BertForTokenClassification.from_pretrained(args.bert_model, num_labels=num_labels)
model.to(device)
if args.do_eval and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
eval_examples = read_ner_example(args.dev_file, args.do_lower_case)
eval_features = convert_examples_to_features(
eval_examples, label_list, args.max_seq_length, tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features], dtype=torch.long)
all_label_masks = torch.tensor([f.label_mask for f in eval_features], dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids, all_label_masks)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
pred_list = []
label_list = []
for input_ids, input_mask, segment_ids, label_ids, label_masks in tqdm(eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
label_masks = label_masks.to(device)
with torch.no_grad():
tmp_eval_loss = model(input_ids, segment_ids, input_mask, label_ids, label_masks)
logits = model(input_ids, segment_ids, input_mask)
active_loss = label_masks.view(-1) == 1
active_logits = logits.view(-1, num_labels)[active_loss]
#print(active_logits.shape)
active_labels = label_ids.view(-1)[active_loss]
active_logits = active_logits.detach().cpu().numpy()
#print(active_logits.shape)
active_labels = active_labels.to('cpu').numpy()
active_preds = np.argmax(active_logits, axis=1)
#print(active_labels.shape, active_preds.shape)
#tmp_eval_accuracy = accuracy(logits, label_ids, label_masks)
#eval_loss += tmp_eval_loss.mean().item()
#eval_accuracy += tmp_eval_accuracy
pred_list.extend(active_preds)
label_list.extend(active_labels)
#print(active_labels.shape)
nb_eval_examples += active_labels.shape[0]
nb_eval_steps += 1
#eval_loss = eval_loss / nb_eval_steps
#eval_accuracy = eval_accuracy / nb_eval_examples
loss = tr_loss/nb_tr_steps if args.do_train else None
eval_f1_micro = f1_score(label_list, pred_list, average='micro')
eval_f1_none = f1_score(label_list, pred_list, average=None)
result = {'eval_f1_micro': eval_f1_micro,
'eval_f1_none': eval_f1_none,
'global_step': global_step,
'loss': loss}
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
output_pred_file = os.path.join(args.output_dir, "pred_results.conll")
label_map = get_labels()
print(len(label_list), len(pred_list))
with open(output_pred_file, 'w') as f, open(args.dev_file) as dev_f:
idx = 1
for l, p, dl in zip(label_list, pred_list, dev_f):
if len(dl) == 0:
print(dl)
f.write('\n')
idx = 1
continue
f.write(' '.join((str(idx), label_map[l], label_map[p])) + '\n')
idx += 1
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help="The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument(
"--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help="The output directory where the model predictions and checkpoints will be written."
)
# Other parameters
parser.add_argument(
"--cache_dir",
default="",
type=str,
help="Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument(
"--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument(
"--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument(
"--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument(
"--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument(
"--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument(
"--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument(
"--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument(
"--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument(
"--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument(
'--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument(
'--server_ip',
type=str,
default='',
help="Can be used for distant debugging.")
parser.add_argument(
'--server_port',
type=str,
default='',
help="Can be used for distant debugging.")
args = parser.parse_args()
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
processors = {"ner": NerProcessor}
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
label_list = processor.get_labels()
num_labels = len(label_list) + 1
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
train_examples = None
num_train_optimization_steps = None
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
print("num_train_optimization_steps: ", num_train_optimization_steps)
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(
global_step / num_train_optimization_steps,
args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
label_map = {i: label for i, label in enumerate(label_list, 1)}
model_config = {
"bert_model": args.bert_model,
"do_lower": args.do_lower_case,
"max_seq_length": args.max_seq_length,
"num_labels": len(label_list) + 1,
"label_map": label_map
}
json.dump(
model_config,
open(os.path.join(args.output_dir, "model_config.json"), "w"))
# Load a trained model and config that you have fine-tuned
else:
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
model.to(device)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
eval_examples = processor.get_dev_examples(args.data_dir)
eval_features = convert_examples_to_features(eval_examples, label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
y_true = []
y_pred = []
label_map = {i: label for i, label in enumerate(label_list, 1)}
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask)
logits = torch.argmax(F.log_softmax(logits, dim=2), dim=2)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
input_mask = input_mask.to('cpu').numpy()
for i, mask in enumerate(input_mask):
temp_1 = []
temp_2 = []
for j, m in enumerate(mask):
if j == 0:
continue
if m:
if label_map[label_ids[i][j]] != "X":
temp_1.append(label_map[label_ids[i][j]])
temp_2.append(label_map[logits[i][j]])
else:
temp_1.pop()
temp_2.pop()
break
y_true.append(temp_1)
y_pred.append(temp_2)
report = classification_report(y_true, y_pred, digits=4)
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
logger.info("\n%s", report)
writer.write(report)
def train_and_evaluate(OUTPUT_DIR, do_train=True, do_eval=True):
""" Train and evaluate a BERT NER Model"""
BATCH_SIZE = 32
LEARNING_RATE = 2e-5
NUM_TRAIN_EPOCHS = 5.0
#in this steps lr will be low and training will be slow
WARMUP_PROPORTION = 0.1
if os.path.exists(OUTPUT_DIR) and os.listdir(OUTPUT_DIR) and do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
OUTPUT_DIR))
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased",
do_lower_case=True)
if do_train:
train_examples, num_train_examples = create_datasets("AGE/train.txt")
num_train_steps = int(
math.ceil(num_train_examples / BATCH_SIZE * NUM_TRAIN_EPOCHS))
num_warmup_steps = int(num_train_steps * WARMUP_PROPORTION)
model = BertForTokenClassification.from_pretrained(
"bert-base-uncased", num_labels=num_labels)
model.to(device)
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in param_optimizer
if not any(nd in n for nd in no_decay)
],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
optimizer = BertAdam(optimizer_grouped_parameters,
lr=LEARNING_RATE,
warmup=WARMUP_PROPORTION,
t_total=num_train_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
train_features = convert_examples_to_features(train_examples,
label_list,
MAX_SEQ_LENGTH,
tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", num_train_examples)
logger.info(" Batch size = %d", BATCH_SIZE)
logger.info(" Num steps = %d", num_train_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=BATCH_SIZE)
model.train()
# for name, param in model.named_parameters():
# if param.requires_grad:
# print(name)
# return
for _ in trange(int(NUM_TRAIN_EPOCHS), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
optimizer.step()
optimizer.zero_grad()
global_step += 1
print(tr_loss)
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(OUTPUT_DIR, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(OUTPUT_DIR, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
label_map = {i: label for i, label in enumerate(label_list, 1)}
model_config = {
"bert_model": "bert-base-uncased",
"do_lower": True,
"max_seq_length": MAX_SEQ_LENGTH,
"num_labels": len(label_list) + 1,
"label_map": label_map
}
json.dump(model_config,
open(os.path.join(OUTPUT_DIR, "model_config.json"), "w"))
else:
output_config_file = os.path.join(OUTPUT_DIR, CONFIG_NAME)
output_model_file = os.path.join(OUTPUT_DIR, WEIGHTS_NAME)
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
model.to(device)
if do_eval:
EVAL_BATCH_SIZE = 32
eval_examples, num_eval_examples = create_datasets("AGE/valid.txt")
eval_features = convert_examples_to_features(eval_examples, label_list,
MAX_SEQ_LENGTH, tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", num_eval_examples)
logger.info(" Batch size = %d", EVAL_BATCH_SIZE)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# # Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=EVAL_BATCH_SIZE)
model.eval()
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
y_true = []
y_pred = []
label_map = {i: label for i, label in enumerate(label_list, 1)}
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask)
logits = torch.argmax(F.log_softmax(logits, dim=2), dim=2)
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
input_mask = input_mask.to('cpu').numpy()
for i, mask in enumerate(input_mask):
temp_1 = []
temp_2 = []
for j, m in enumerate(mask):
if j == 0:
continue
if m:
if label_map[label_ids[i][j]] != "X":
temp_1.append(label_map[label_ids[i][j]])
temp_2.append(label_map[logits[i][j]])
else:
temp_1.pop()
temp_2.pop()
break
y_true.append(temp_1)
y_pred.append(temp_2)
report = classification_report(y_true, y_pred)
output_eval_file = os.path.join(OUTPUT_DIR, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
logger.info("\n%s", report)
writer.write(report)
max_seq_length = 512
model_path = "/.pytorch_pretrained_bert/token_model.pt"
bert_model = "/.pytorch_pretrained_bert/bert-base-uncased.tar.gz"
bert_vocab = "/.pytorch_pretrained_bert/bert-base-uncased-vocab.txt"
# device = torch.device("cuda:2" if torch.cuda.is_available() else "cpu")
device = "cpu"
tokenizer = BertTokenizer.from_pretrained(bert_vocab)
label_list = [
"B-etime", "B-fname", "B-organizer", "B-participant", "B-place",
"B-target", "B-trigger", "I-etime", "I-fname", "I-organizer",
"I-participant", "I-place", "I-target", "I-trigger", "O"
]
label_map = {}
for (i, label) in enumerate(label_list):
label_map[i] = label
model = BertForTokenClassification.from_pretrained(
bert_model, PYTORCH_PRETRAINED_BERT_CACHE, num_labels=len(label_list))
if device == "cpu":
model.load_state_dict(torch.load(model_path, map_location='cpu'))
else:
model.load_state_dict(torch.load(model_path))
model.to(device)
api.add_resource(queryList, '/queries')
app.run(host='0.0.0.0', port=4998, debug=True)
def main():
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
## Other parameters
parser.add_argument(
"--cache_dir",
default="",
type=str,
help=
"Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument(
"--max_seq_length",
default=128,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help=
"Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument('--server_ip',
type=str,
default='',
help="Can be used for distant debugging.")
parser.add_argument('--server_port',
type=str,
default='',
help="Can be used for distant debugging.")
args = parser.parse_args()
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logging.basicConfig(
format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
labels = list('BIEOS')
processor = sequence_labeling.NerProcessor(labels)
label_list = processor.labels
num_labels = len(label_list)
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
train_examples = None
num_train_optimization_steps = None
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
if args.do_train:
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in param_optimizer
if not any(nd in n for nd in no_decay)
],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
warmup_linear = WarmupLinearSchedule(
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
train_features = processor.convert_examples_to_features(
train_examples, args.max_seq_length, tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_ids for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
# define a new function to compute loss values for both output_modes
logits = model(input_ids, segment_ids, input_mask, labels=None)
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits.view(-1, num_labels),
label_ids.view(-1))
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear.get_lr(
global_step, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
if args.do_train and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
# Save a trained model, configuration and tokenizer
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(args.output_dir)
processor.save(
os.path.join(args.output_dir, sequence_labeling.PROCESSOR_NAME))
# Load a trained model and vocabulary that you have fine-tuned
model = BertForTokenClassification.from_pretrained(
args.output_dir, num_labels=num_labels)
tokenizer = BertTokenizer.from_pretrained(
args.output_dir, do_lower_case=args.do_lower_case)
else:
model = BertForTokenClassification.from_pretrained(
args.bert_model, num_labels=num_labels)
model.to(device)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
eval_examples = processor.get_dev_examples(args.data_dir)
eval_features = processor.convert_examples_to_features(
eval_examples, args.max_seq_length, tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_ids for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss = 0
nb_eval_steps = 0
preds = []
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask, labels=None)
# create eval loss and other metric required by the task
loss_fct = CrossEntropyLoss()
tmp_eval_loss = loss_fct(logits.view(-1, num_labels),
label_ids.view(-1))
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if len(preds) == 0:
preds.append(logits.detach().cpu().numpy())
else:
preds[0] = np.append(preds[0],
logits.detach().cpu().numpy(),
axis=0)
eval_loss = eval_loss / nb_eval_steps
preds = preds[0]
preds = np.argmax(preds, axis=2)
print(preds.shape)
print(all_label_ids.numpy().shape)
result = compute_metrics(preds.flatten(),
all_label_ids.numpy().flatten())
loss = tr_loss / global_step if args.do_train else None
result['eval_loss'] = eval_loss
result['global_step'] = global_step
result['loss'] = loss
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
def model_train(bert_model,
max_seq_length,
do_lower_case,
num_train_epochs,
train_batch_size,
gradient_accumulation_steps,
learning_rate,
weight_decay,
loss_scale,
warmup_proportion,
processor,
device,
n_gpu,
fp16,
cache_dir,
local_rank,
dry_run,
no_cuda,
output_dir=None):
label_map = processor.get_labels()
if gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(gradient_accumulation_steps))
train_batch_size = train_batch_size // gradient_accumulation_steps
train_dataloader = processor.get_train_examples(train_batch_size,
local_rank)
# Batch sampler divides by batch_size!
num_train_optimization_steps = int(
len(train_dataloader) * num_train_epochs / gradient_accumulation_steps)
if local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
cache_dir = cache_dir if cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE),
'distributed_{}'.format(local_rank))
model = BertForTokenClassification.from_pretrained(
bert_model, cache_dir=cache_dir, num_labels=len(label_map))
if fp16:
model.half()
model.to(device)
if local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
weight_decay
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer, static_loss_scale=loss_scale)
warmup_linear = WarmupLinearSchedule(
warmup=warmup_proportion, t_total=num_train_optimization_steps)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=learning_rate,
warmup=warmup_proportion,
t_total=num_train_optimization_steps)
warmup_linear = None
global_step = 0
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataloader))
logger.info(" Batch size = %d", train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
logger.info(" Num epochs = %d", num_train_epochs)
model_config = {
"bert_model": bert_model,
"do_lower": do_lower_case,
"max_seq_length": max_seq_length,
"label_map": label_map
}
def save_model(lh):
if output_dir is None:
return
output_model_file = os.path.join(output_dir,
"pytorch_model_ep{}.bin".format(ep))
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(output_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
json.dump(model_config,
open(os.path.join(output_dir, "model_config.json"), "w"))
lh = pd.DataFrame(lh, columns=['global_step', 'loss'])
loss_history_file = os.path.join(output_dir,
"loss_ep{}.pkl".format(ep))
lh.to_pickle(loss_history_file)
def load_model(epoch):
if output_dir is None:
return False
output_model_file = os.path.join(
output_dir, "pytorch_model_ep{}.bin".format(epoch))
if not os.path.exists(output_model_file):
return False
logger.info("Loading epoch {} from disk...".format(epoch))
model.load_state_dict(
torch.load(output_model_file,
map_location=lambda storage, loc: storage
if no_cuda else None))
return True
model.train()
for ep in trange(1, int(num_train_epochs) + 1, desc="Epoch"):
if dry_run and ep > 1:
logger.info("Dry run. Stop.")
break
if load_model(ep):
global_step += len(train_dataloader) // gradient_accumulation_steps
continue
loss_history = list()
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
with tqdm(total=len(train_dataloader), desc=f"Epoch {ep}") as pbar:
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if gradient_accumulation_steps > 1:
loss = loss / gradient_accumulation_steps
if fp16:
optimizer.backward(loss)
else:
loss.backward()
loss_history.append((global_step, loss.item()))
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
pbar.update(1)
mean_loss = tr_loss * gradient_accumulation_steps / nb_tr_steps
pbar.set_postfix_str(f"Loss: {mean_loss:.5f}")
if dry_run and len(loss_history) > 2:
logger.info("Dry run. Stop.")
break
if (step + 1) % gradient_accumulation_steps == 0:
if fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = learning_rate * warmup_linear.get_lr(
global_step, warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
save_model(loss_history)
return model, model_config
num_workers=4,
collate_fn=NerDataset.pad)
#%%
'''
##### Use only BertForTokenClassification #####
'''
print('*** Use only BertForTokenClassification ***')
if load_checkpoint and os.path.exists(output_dir+'/ner_bert_checkpoint.pt'):
checkpoint = torch.load(output_dir+'/ner_bert_checkpoint.pt', map_location='cpu')
start_epoch = checkpoint['epoch']+1
valid_acc_prev = checkpoint['valid_acc']
valid_f1_prev = checkpoint['valid_f1']
model = BertForTokenClassification.from_pretrained(
bert_model_scale, state_dict=checkpoint['model_state'], num_labels=len(label_list))
print('Loaded the pretrain NER_BERT model, epoch:',checkpoint['epoch'],'valid acc:',
checkpoint['valid_acc'], 'valid f1:', checkpoint['valid_f1'])
else:
start_epoch = 0
valid_acc_prev = 0
valid_f1_prev = 0
model = BertForTokenClassification.from_pretrained(
bert_model_scale, num_labels=len(label_list))
model.to(device)
# Prepare optimizer
named_params = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument("--data_dir",
default=None,
type=str,
required=True,
help="The input data dir.")
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help=
"Bert pre-trained model selected in the list: bert-base-uncased, bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, bert-base-multilingual-cased, bert-base-chinese."
)
parser.add_argument("--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
# Other parameters
parser.add_argument(
"--max_seq_length",
default=128,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization. \nSequences longer than this will be truncated, and sequences shorter \nthan this will be padded."
)
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training."
)
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help=
"Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n0 (default value): dynamic loss scaling.\nPositive power of 2: static loss scaling value.\n"
)
args = parser.parse_args()
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1 # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
os.makedirs(os.path.join(args.output_dir, args.task_name), exist_ok=True)
processor = DataProcessor(os.path.join(args.data_dir, args.task_name),
do_lower_case=args.do_lower_case)
train_examples = processor.get_train_examples()
label_list = processor.all_labels
num_labels = len(label_list)
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
num_train_optimization_steps = None
if args.do_train:
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
model = BertForTokenClassification.from_pretrained(args.bert_model,
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex to use distributed and fp16 training.")
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
tokenizer)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_ids for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_steps = 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses if args.fp16 is False,
# BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(
global_step / num_train_optimization_steps,
args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, args.task_name,
"pytorch_model.bin")
if args.do_train:
torch.save(model_to_save.state_dict(), output_model_file)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
# Load a trained model that you have fine-tuned
model_state_dict = torch.load(output_model_file)
model = BertForTokenClassification.from_pretrained(
args.bert_model,
state_dict=model_state_dict,
num_labels=num_labels)
model.to(device)
eval_examples = processor.get_dev_examples()
eval_features = convert_examples_to_features(eval_examples, label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_ids for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
all_examples = []
eval_loss, eval_accuracy = 0, 0
nb_eval_steps, nb_eval_examples = 0, 0
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
tmp_eval_loss = model(input_ids, segment_ids, input_mask,
label_ids)
logits = model(input_ids, segment_ids, input_mask)
examples = get_output_file(logits, input_ids, input_mask,
label_ids, tokenizer.ids_to_tokens,
processor.all_labels)
all_examples.extend(examples)
eval_loss += tmp_eval_loss.mean().item()
nb_eval_examples += input_ids.size(0)
nb_eval_steps += 1
eval_loss = eval_loss / nb_eval_steps
acc, p, r, f1 = get_ner_fmeasure([e.labels for e in all_examples],
[e.predicts for e in all_examples])
result = {
'eval_loss': eval_loss,
'eval_accuracy': acc,
'precision': p,
'recall': r,
'f_meature': f1,
}
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
output_eval_file = os.path.join(args.output_dir, args.task_name,
"eval_results.txt")
with open(output_eval_file, "w") as writer:
for example in all_examples:
for token, label, pred in zip(example.text, example.labels,
example.predicts):
writer.write(F"{token} {label} {pred}\n")
writer.write('\n')
def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help="The input data dir. Should contain the .tsv files (or other data files) for the task."
)
parser.add_argument(
"--log_dir",
default=None,
type=str,
required=True,
help="The log dir. Should contain the .txt file (or other data file) for the task."
)
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
parser.add_argument("--task_name",
default=None,
type=str,
required=True,
help="The name of the task to train.")
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help="The output directory where the model predictions and checkpoints will be written."
)
# Other parameters
parser.add_argument(
"--cache_dir",
default="",
type=str,
help="Where do you want to store the pre-trained models downloaded from s3")
parser.add_argument(
"--max_seq_length",
default=128,
type=int,
help="The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help="Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help="Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
parser.add_argument('--server_ip',
type=str,
default='',
help="Can be used for distant debugging.")
parser.add_argument('--server_port',
type=str,
default='',
help="Can be used for distant debugging.")
args = parser.parse_args()
# log setting
handler = logging.FileHandler(os.path.join(args.log_dir, "log.txt"))
handler.setFormatter(logging.DEBUG)
formatter = logging.Formatter(
'%(asctime)s - %(levelname)s - %(name)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
if args.server_ip and args.server_port:
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
import ptvsd
print("Waiting for debugger attach")
ptvsd.enable_attach(address=(args.server_ip, args.server_port),
redirect_output=True)
ptvsd.wait_for_attach()
processors = {
# bbn processor
"bbn": BBNNerProcessor,
}
output_modes = {
"bbn": "classification",
}
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
args.train_batch_size = args.train_batch_size // args.gradient_accumulation_steps
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
task_name = args.task_name.lower()
if task_name not in processors:
raise ValueError("Task not found: %s" % (task_name))
processor = processors[task_name]()
output_mode = output_modes[task_name]
if task_name == 'bbn':
label_list = processor.get_labels(args.data_dir)
else:
label_list = processor.get_labels()
num_labels = len(label_list)
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
train_examples = None
num_train_optimization_steps = None
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_optimization_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps) * args.num_train_epochs
if args.local_rank != -1:
num_train_optimization_steps = num_train_optimization_steps // torch.distributed.get_world_size(
)
# Prepare model
cache_dir = args.cache_dir if args.cache_dir else os.path.join(
str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed_{}'.format(
args.local_rank))
model = BertForTokenClassification.from_pretrained(args.bert_model,
cache_dir=cache_dir,
num_labels=num_labels)
if args.fp16:
model.half()
try:
model.to(device)
except Exception:
logger.warning("toGPU failed, failed msg:" + traceback.format_exc())
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
# prepare Data
# train_label_ids, dev_label_ids, test_label_ids = process_data(tokenizer, processor, args.data_dir, args.max_seq_length)
if args.do_train:
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_optimization_steps)
train_data = torch.load(os.path.join(args.data_dir, "train.pt"))
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in range(int(args.num_train_epochs)):
tr_loss = 0
last_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(train_dataloader):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
# define a new function to compute loss values for both output_modes
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear(
global_step / num_train_optimization_steps,
args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
if abs(loss.item() - last_loss) <= 5e-10:
break
# if abs(loss.item() - last_loss) != 0:
# print("iterate fine")
# print("step: " + str(step))
# print(abs(loss.item() - last_loss))
last_loss = loss.item()
# Save a trained model and the associated configuration
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
torch.save(model_to_save.state_dict(), output_model_file)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
with open(output_config_file, 'w') as f:
f.write(model_to_save.config.to_json_string())
# Load a trained model and config that you have fine-tuned
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
else:
logger.info("preparing model")
# model = BertForTokenClassification.from_pretrained(
# args.bert_model, num_labels=num_labels)
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
config = BertConfig(output_config_file)
model = BertForTokenClassification(config, num_labels=num_labels)
model.load_state_dict(torch.load(output_model_file))
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
model.to(device)
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
model.eval()
# eval_examples = processor.get_dev_examples(args.data_dir)
eval_examples = processor.get_dev_examples(args.data_dir)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
eval_data = torch.load(os.path.join(args.data_dir, "dev.pt"))
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss = 0
nb_eval_steps = 0
preds = []
active_labels_dataset = []
i = 0
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits, active_loss = model(input_ids, segment_ids, input_mask, labels=None)
active_labels = label_ids.view(-1)[active_loss]
# create eval loss and other metric required by the task
if output_mode == "classification":
loss_fct = CrossEntropyLoss()
# tmp_eval_loss = loss_fct(logits.view(-1, num_labels),
# active_labels.view(-1))
tmp_eval_loss = 0
elif output_mode == "regression":
loss_fct = MSELoss()
tmp_eval_loss = loss_fct(logits.view(-1), active_labels.view(-1))
# eval_loss += tmp_eval_loss.mean().item()
eval_loss += 0
nb_eval_steps += 1
# if len(preds) == 0:
# preds.append(logits.detach().cpu().numpy())
# else:
# preds[0] = np.append(preds[0],
# logits.detach().cpu().numpy(),
# axis=0)
logits = np.argmax(logits.detach().cpu().numpy(), axis=1)
preds.append(logits)
active_labels_dataset.append(active_labels)
eval_loss = eval_loss / nb_eval_steps
# preds = preds[0]
preds_flat = []
labels_flat = []
for s in preds:
for l in s: # l is label
preds_flat.append(l)
for s in active_labels_dataset:
for l in s:
labels_flat.append(l.detach().cpu().numpy())
preds_flat = np.array(preds_flat)
labels_flat = np.array(labels_flat)
for i in range(len(preds_flat)):
if preds_flat[i] == 37:
preds_flat[i] = 7
elif preds_flat[i] == 34:
preds_flat[i] == 12
elif preds_flat[i] == 26:
preds_flat[i] = 36
elif preds_flat[i] == 36:
preds_flat[i] = 37
elif preds_flat[i] == 41:
preds_flat[i] = 34
elif preds_flat[i] == 31:
preds_flat[i] = 39
elif preds_flat[i] == 15:
preds_flat[i] = 38
# label_map = dict()
# for i in range(len(preds_flat)):
# key = str(preds_flat[i]) + '-' + str(labels_flat[i])
# if key in label_map.keys():
# label_map[key] += 1
# else:
# label_map[key] = 0
# for k in label_map.keys():
# if label_map[k] > 1000:
# print(k, ":", label_map[k])
# if output_mode == "classification":
# preds = np.argmax(preds, axis=2)
# elif output_mode == "regression":
# preds = np.squeeze(preds)
result = compute_metrics(task_name, preds_flat, labels_flat)
loss = tr_loss / nb_tr_steps if args.do_train else None
result['eval_loss'] = eval_loss
result['global_step'] = global_step
result['loss'] = loss
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
# hack for MNLI-MM
if task_name == "mnli":
task_name = "mnli-mm"
processor = processors[task_name]()
if os.path.exists(args.output_dir +
'-MM') and os.listdir(args.output_dir +
'-MM') and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".
format(args.output_dir))
if not os.path.exists(args.output_dir + '-MM'):
os.makedirs(args.output_dir + '-MM')
eval_examples = processor.get_dev_examples(args.data_dir)
eval_features = convert_examples_to_features(
eval_examples, label_list, args.max_seq_length, tokenizer,
output_mode)[0]
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor(
[f.input_mask for f in eval_features], dtype=torch.long)
all_segment_ids = torch.tensor(
[f.segment_ids for f in eval_features], dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
eval_loss = 0
nb_eval_steps = 0
preds = []
for input_ids, input_mask, segment_ids, label_ids in eval_dataloader:
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
with torch.no_grad():
logits = model(input_ids,
segment_ids,
input_mask,
labels=None)
loss_fct = CrossEntropyLoss()
tmp_eval_loss = loss_fct(logits.view(-1, num_labels),
label_ids.view(-1))
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if len(preds) == 0:
preds.append(logits.detach().cpu().numpy())
else:
preds[0] = np.append(preds[0],
logits.detach().cpu().numpy(),
axis=0)
eval_loss = eval_loss / nb_eval_steps
preds = preds[0]
preds = np.argmax(preds, axis=1)
result = compute_metrics(task_name, preds, all_label_ids.numpy())
loss = tr_loss / nb_tr_steps if args.do_train else None
result['eval_loss'] = eval_loss
result['global_step'] = global_step
result['loss'] = loss
output_eval_file = os.path.join(args.output_dir + '-MM',
"eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
def main():
"""Main method of this module."""
parser = argparse.ArgumentParser()
parser.add_argument("-c",
"--input_file",
default=None,
type=str,
required=True,
help="The input data dir")
parser.add_argument("-o",
"--output_file",
default=None,
type=str,
required=True,
help="Output file for predictions")
parser.add_argument("--bert_model",
default="",
type=str,
required=True,
help="Bert pre-trained model path")
parser.add_argument("--bert_tokenizer",
default="",
type=str,
required=True,
help="Bert tokenizer path")
parser.add_argument("--model_load",
default="",
type=str,
required=True,
help="The path of model state.")
parser.add_argument(
"--max_seq_length",
default=512,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
args = parser.parse_args()
max_seq_length = args.max_seq_length
model_path = args.model_load
input_file = args.input_file
output_file = args.output_file
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BertTokenizer.from_pretrained(args.bert_tokenizer)
label_list = [
"B-etime", "B-fname", "B-loc", "B-organizer", "B-participant",
"B-place", "B-target", "B-trigger", "I-etime", "I-fname", "I-loc",
"I-organizer", "I-participant", "I-place", "I-target", "I-trigger", "O"
]
model = BertForTokenClassification.from_pretrained(
args.bert_model,
PYTORCH_PRETRAINED_BERT_CACHE,
num_labels=len(label_list))
label_map = {}
for (i, label) in enumerate(label_list):
label_map[i] = label
# try:
# model.load_state_dict(torch.load(model_path)) # , map_location='cpu' for only cpu
# except: #When model is parallel
# model = torch.nn.DataParallel(model)
# model.load_state_dict(torch.load(model_path)) # , map_location='cpu' for only cpu
model.load_state_dict(torch.load(model_path))
logger.info("Model state has been loaded.")
model.to(device)
with open(input_file, "r", encoding="utf-8") as f:
lines = f.read().splitlines()
examples = []
words = []
for (i, line) in enumerate(lines):
line = line.strip()
if line == "SAMPLE_START":
words.append("[CLS]")
elif line == "[SEP]":
continue
elif line == "":
tokens = []
for (j, word) in enumerate(words):
if word == "[CLS]":
tokens.append("[CLS]")
continue
tokenized = tokenizer.tokenize(word)
tokens.append(tokenized[0])
if len(tokens) > max_seq_length - 1:
tokens = tokens[0:(max_seq_length - 1)]
tokens.append("[SEP]")
tokens = tokenizer.convert_tokens_to_ids(tokens)
segment_ids = [0] * len(tokens)
input_mask = [1] * len(tokens)
while len(tokens) < max_seq_length:
tokens.append(0)
segment_ids.append(0)
input_mask.append(0)
examples.append((tokens, input_mask, segment_ids))
words = []
continue
elif line in ["\x91", "\x92", "\x97"]:
continue
else:
words.append(line)
# print(examples)
all_labels = []
model.eval()
for (input_ids, input_mask, segment_ids) in examples:
org_input_mask = input_mask
org_input_mask = [x for x in org_input_mask if x != 0]
input_ids = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
input_mask = torch.tensor(input_mask, dtype=torch.long).unsqueeze(0)
segment_ids = torch.tensor(segment_ids, dtype=torch.long).unsqueeze(0)
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
with torch.no_grad():
logits = model(input_ids, segment_ids, input_mask)
logits = logits.detach().cpu().numpy()
labels = np.argmax(logits, axis=-1).reshape(-1)
labels = labels[0:len(org_input_mask)]
# while len(labels) < max_seq_length:
# labels = np.append(labels, 16) # Add "O"
all_labels = np.append(all_labels, labels)
j = 0
count = 0
with open(output_file, "w", encoding="utf-8") as g:
for (i, line) in enumerate(lines):
line = line.strip()
if line == "SAMPLE_START":
count += 1
g.write("SAMPLE_START\tO\n")
j += 1
elif line == "[SEP]":
g.write("[SEP]\tO\n")
elif line == "\x91":
g.write("\x91\tO\n")
elif line == "\x92":
g.write("\x92\tO\n")
elif line == "\x97":
g.write("\x97\tO\n")
elif line == "":
g.write("\n")
count = 0
j += 1 # We have a SEP at the end
else:
count += 1
if count < max_seq_length:
g.write(line + "\t" + label_map[all_labels[j]] + "\n")
j += 1
else:
g.write(line + "\tO\n")
logger.info("The predictions have been written to the output folder.")
def main():
parser = argparse.ArgumentParser()
## Required parameters
# Data Directory
parser.add_argument(
"--data_dir",
default=None,
type=str,
required=True,
help=
"The input data dir. Should contain the .tsv files (or other data files) for the task."
)
# Bert Model
parser.add_argument(
"--bert_model",
default=None,
type=str,
required=True,
help="Bert pre-trained model selected in the list: bert-base-uncased, "
"bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, "
"bert-base-multilingual-cased, bert-base-chinese.")
# Output Directory
parser.add_argument(
"--output_dir",
default=None,
type=str,
required=True,
help=
"The output directory where the model predictions and checkpoints will be written."
)
## Other parameters
# Max sequence length
parser.add_argument(
"--max_seq_length",
default=128,
type=int,
help=
"The maximum total input sequence length after WordPiece tokenization. \n"
"Sequences longer than this will be truncated, and sequences shorter \n"
"than this will be padded.")
# Train it?
parser.add_argument("--do_train",
action='store_true',
help="Whether to run training.")
# Run evaluation?
parser.add_argument("--do_eval",
action='store_true',
help="Whether to run eval on the dev set.")
# Uncased?
parser.add_argument(
"--do_lower_case",
action='store_true',
help="Set this flag if you are using an uncased model.")
# Set batch size
parser.add_argument("--train_batch_size",
default=32,
type=int,
help="Total batch size for training.")
# Batch size for evaluation
parser.add_argument("--eval_batch_size",
default=8,
type=int,
help="Total batch size for eval.")
# Learning Rate for Adam
parser.add_argument("--learning_rate",
default=5e-5,
type=float,
help="The initial learning rate for Adam.")
# Training epochs
parser.add_argument("--num_train_epochs",
default=3.0,
type=float,
help="Total number of training epochs to perform.")
# ??
parser.add_argument(
"--warmup_proportion",
default=0.1,
type=float,
help=
"Proportion of training to perform linear learning rate warmup for. "
"E.g., 0.1 = 10%% of training.")
parser.add_argument("--no_cuda",
action='store_true',
help="Whether not to use CUDA when available")
parser.add_argument("--local_rank",
type=int,
default=-1,
help="local_rank for distributed training on gpus")
parser.add_argument('--seed',
type=int,
default=42,
help="random seed for initialization")
parser.add_argument(
'--gradient_accumulation_steps',
type=int,
default=1,
help=
"Number of updates steps to accumulate before performing a backward/update pass."
)
parser.add_argument(
'--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument(
'--loss_scale',
type=float,
default=0,
help=
"Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
args = parser.parse_args()
# cuda or cpu
if args.local_rank == -1 or args.no_cuda:
device = torch.device("cuda" if torch.cuda.is_available()
and not args.no_cuda else "cpu")
n_gpu = torch.cuda.device_count()
else:
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
n_gpu = 1
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
logger.info(
"device: {} n_gpu: {}, distributed training: {}, 16-bits training: {}".
format(device, n_gpu, bool(args.local_rank != -1), args.fp16))
# Check for valid args
if args.gradient_accumulation_steps < 1:
raise ValueError(
"Invalid gradient_accumulation_steps parameter: {}, should be >= 1"
.format(args.gradient_accumulation_steps))
# Set train batch size
args.train_batch_size = int(args.train_batch_size /
args.gradient_accumulation_steps)
# Seeds
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(args.seed)
if not args.do_train and not args.do_eval:
raise ValueError(
"At least one of `do_train` or `do_eval` must be True.")
if os.path.exists(args.output_dir) and os.listdir(
args.output_dir) and args.do_train:
raise ValueError(
"Output directory ({}) already exists and is not empty.".format(
args.output_dir))
os.makedirs(args.output_dir, exist_ok=True)
processor = NerProcessor()
label_list = processor.get_labels()
num_labels = len(label_list)
# Call Tokenizer
tokenizer = BertTokenizer.from_pretrained(args.bert_model,
do_lower_case=args.do_lower_case)
train_examples = None
num_train_steps = None
if args.do_train:
train_examples = processor.get_train_examples(args.data_dir)
num_train_steps = int(
len(train_examples) / args.train_batch_size /
args.gradient_accumulation_steps * args.num_train_epochs)
# Can I use Token Classification here?
model = BertForTokenClassification.from_pretrained(
args.bert_model,
cache_dir=PYTORCH_PRETRAINED_BERT_CACHE /
'distributed_{}'.format(args.local_rank),
num_labels=num_labels)
if args.fp16:
model.half()
model.to(device)
if args.local_rank != -1:
try:
from apex.parallel import DistributedDataParallel as DDP
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
model = DDP(model)
elif n_gpu > 1:
model = torch.nn.DataParallel(model)
# Prepare optimizer
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params':
[p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay':
0.01
}, {
'params':
[p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay':
0.0
}]
t_total = num_train_steps
if args.local_rank != -1:
t_total = t_total // torch.distributed.get_world_size()
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError(
"Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training."
)
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer,
static_loss_scale=args.loss_scale)
else:
# Create Optimizer
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=t_total)
global_step = 0
nb_tr_steps = 0
tr_loss = 0
if args.do_train:
train_features = convert_examples_to_features(train_examples,
label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", args.train_batch_size)
logger.info(" Num steps = %d", num_train_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in train_features],
dtype=torch.long)
train_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
if args.local_rank == -1:
train_sampler = RandomSampler(train_data)
else:
train_sampler = DistributedSampler(train_data)
train_dataloader = DataLoader(train_data,
sampler=train_sampler,
batch_size=args.train_batch_size)
model.train()
for _ in trange(int(args.num_train_epochs), desc="Epoch"):
tr_loss = 0
nb_tr_examples, nb_tr_steps = 0, 0
for step, batch in enumerate(
tqdm(train_dataloader, desc="Iteration")):
batch = tuple(t.to(device) for t in batch)
input_ids, input_mask, segment_ids, label_ids = batch
loss = model(input_ids, segment_ids, input_mask, label_ids)
if n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu.
if args.gradient_accumulation_steps > 1:
loss = loss / args.gradient_accumulation_steps
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
tr_loss += loss.item()
nb_tr_examples += input_ids.size(0)
nb_tr_steps += 1
if (step + 1) % args.gradient_accumulation_steps == 0:
# modify learning rate with special warm up BERT uses
lr_this_step = args.learning_rate * warmup_linear(
global_step / t_total, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
global_step += 1
# Save a trained model
model_to_save = model.module if hasattr(
model, 'module') else model # Only save the model it-self
output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
if args.do_train:
torch.save(model_to_save.state_dict(), output_model_file)
# Load a trained model that you have fine-tuned
model_state_dict = torch.load(output_model_file)
model = BertForTokenClassification.from_pretrained(
args.bert_model, state_dict=model_state_dict, num_labels=num_labels)
model.to(device)
######################################################################
outputwrite = open('debug', 'w')
######################################################################
if args.do_eval and (args.local_rank == -1
or torch.distributed.get_rank() == 0):
eval_examples = processor.get_test_examples(args.data_dir)
eval_features = convert_examples_to_features(eval_examples, label_list,
args.max_seq_length,
tokenizer)
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", args.eval_batch_size)
all_input_ids = torch.tensor([f.input_ids for f in eval_features],
dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in eval_features],
dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in eval_features],
dtype=torch.long)
all_label_ids = torch.tensor([f.label_id for f in eval_features],
dtype=torch.long)
eval_data = TensorDataset(all_input_ids, all_input_mask,
all_segment_ids, all_label_ids)
# Run prediction for full data
eval_sampler = SequentialSampler(eval_data)
eval_dataloader = DataLoader(eval_data,
sampler=eval_sampler,
batch_size=args.eval_batch_size)
model.eval()
exact_match, num_sentence = 0, 0
tp, tn = 0, 0
fp, fn = 0, 0
for input_ids, input_mask, segment_ids, label_ids in tqdm(
eval_dataloader, desc="Evaluating"):
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)
label_ids = label_ids.to(device)
# print(label_ids)
with torch.no_grad():
tmp_eval_loss = model(input_ids, segment_ids, input_mask,
label_ids)
logits = model(input_ids, segment_ids, input_mask)
### Edited here!!!! - check dimension
try:
for i in range(args.eval_batch_size):
outputwrite.write("Start ----\n")
sentence_flag = True
for j in range(args.max_seq_length):
test_prediction = label_list[int(
torch.argmax(logits[i][j]))]
outputwrite.write(test_prediction + " ")
test_answer = label_list[label_ids[i][j]]
outputwrite.write(test_answer)
outputwrite.write("\n")
if test_prediction == "X":
continue
if test_prediction not in ["O"]: # Positive
if test_answer == test_prediction:
tp += 1
outputwrite.write("True Positive")
else:
fp += 1
outputwrite.write("False Positive")
sentence_flag = False
else:
if test_answer == test_prediction:
tn += 1
outputwrite.write("True Negative")
else:
fn += 1
sentence_flag = False
outputwrite.write("False Negative")
outputwrite.write("\n")
if sentence_flag:
exact_match += 1
outputwrite.write(" - Exact Match - ")
num_sentence += 1
outputwrite.write("\n")
except Exception as e:
pass
logits = logits.detach().cpu().numpy()
label_ids = label_ids.to('cpu').numpy()
exact_match = exact_match / num_sentence
fprecision = tp / (tp + fp)
frecall = tp / (tp + fn)
f1 = 2 * fprecision * frecall / (fprecision + frecall)
result = {
'exact_match': exact_match,
'f1': f1,
'precision': fprecision,
'recall': frecall,
'true_pos': tp,
'true_neg': tn,
'false_pos': fp,
'false_neg': fn
}
output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
logger.info("***** Eval results *****")
for key in sorted(result.keys()):
logger.info(" %s = %s", key, str(result[key]))
writer.write("%s = %s\n" % (key, str(result[key])))
def init_model(self):
self.model = BertForTokenClassification.from_pretrained(
"bert-base-uncased", num_labels=self.num_labels)
self.model.to(self.device)