def run():
"""train the model"""
# set the logger
utils.set_logger(config.log_dir)
logging.info("device: {}".format(config.device))
# 处理数据，分离文本和标签
processor = Processor(config)
processor.process()
logging.info("--------Process Done!--------")
# 分离出验证集
word_train, word_dev, label_train, label_dev = load_dev('train')
# build dataset
train_dataset = NERDataset(word_train, label_train, config)
dev_dataset = NERDataset(word_dev, label_dev, config)
logging.info("--------Dataset Build!--------")
# get dataset size
train_size = len(train_dataset)
# build data_loader
train_loader = DataLoader(train_dataset, batch_size=config.batch_size,
shuffle=True, collate_fn=train_dataset.collate_fn)
dev_loader = DataLoader(dev_dataset, batch_size=config.batch_size,
shuffle=True, collate_fn=dev_dataset.collate_fn)
logging.info("--------Get Dataloader!--------")
# Prepare model
device = config.device
model = BertNER.from_pretrained(config.roberta_model, num_labels=len(config.label2id))
model.to(device)
# Prepare optimizer
if config.full_fine_tuning:
# model.named_parameters(): [bert, classifier, crf]
bert_optimizer = list(model.bert.named_parameters())
classifier_optimizer = list(model.classifier.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in bert_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay': config.weight_decay},
{'params': [p for n, p in bert_optimizer if any(nd in n for nd in no_decay)],
'weight_decay': 0.0},
{'params': [p for n, p in classifier_optimizer if not any(nd in n for nd in no_decay)],
'lr': config.learning_rate * 5, 'weight_decay': config.weight_decay},
{'params': [p for n, p in classifier_optimizer if any(nd in n for nd in no_decay)],
'lr': config.learning_rate * 5, 'weight_decay': 0.0},
{'params': model.crf.parameters(), 'lr': config.learning_rate * 5}
]
# only fine-tune the head classifier
else:
param_optimizer = list(model.classifier.named_parameters())
optimizer_grouped_parameters = [{'params': [p for n, p in param_optimizer]}]
optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate, correct_bias=False)
train_steps_per_epoch = train_size // config.batch_size
scheduler = get_cosine_schedule_with_warmup(optimizer,
num_warmup_steps=(config.epoch_num // 10) * train_steps_per_epoch,
num_training_steps=config.epoch_num * train_steps_per_epoch)
# Train the model
logging.info("--------Start Training!--------")
train(train_loader, dev_loader, model, optimizer, scheduler, config.model_dir)
def pharmacy_counting():
start_time = time.time()
processor = Processor(input_file, output_file)
processor.process()
print("program executed: %s" % (time.time() - start_time))