def main():
parser = argparse.ArgumentParser()
# Required parameters
parser.add_argument("--do_train", default=False, action='store_true')
parser.add_argument('--do_eval', default=False, action='store_true')
parser.add_argument("--do_predict", default=False, action='store_true')
parser.add_argument('--markup',
default='bios',
type=str,
choices=['bios', 'bio'])
parser.add_argument("--arch", default='bilstm_crf', type=str)
parser.add_argument('--learning_rate', default=0.001, type=float)
parser.add_argument('--seed', default=1234, type=int)
# parser.add_argument('--gpu',default='0',type=str)
parser.add_argument('--gpu', default='', type=str)
parser.add_argument('--epochs', default=50, type=int)
parser.add_argument('--batch_size', default=32, type=int)
parser.add_argument('--embedding_size', default=128, type=int)
parser.add_argument('--hidden_size', default=384, type=int)
parser.add_argument("--grad_norm",
default=5.0,
type=float,
help="Max gradient norm.")
parser.add_argument("--task_name", type=str, default='ner')
args = parser.parse_args()
args.data_dir = config.data_dir
if not config.output_dir.exists():
args.output_dir.mkdir()
args.output_dir = config.output_dir / '{}'.format(args.arch)
if not args.output_dir.exists():
args.output_dir.mkdir()
init_logger(log_file=str(args.output_dir /
'{}-{}.log'.format(args.arch, args.task_name)))
seed_everything(args.seed)
if args.gpu != '':
args.device = torch.device(f"cuda:{args.gpu}")
else:
args.device = torch.device("cpu")
args.id2label = {i: label for i, label in enumerate(config.label2id)}
args.label2id = config.label2id
processor = CluenerProcessor(data_dir=config.data_dir)
processor.get_vocab()
model = NERModel(vocab_size=len(processor.vocab),
embedding_size=args.embedding_size,
hidden_size=args.hidden_size,
device=args.device,
label2id=args.label2id)
model.to(args.device)
if args.do_train:
train(args, model, processor)
if args.do_eval:
model_path = args.output_dir / 'best-model.bin'
model = load_model(model, model_path=str(model_path))
evaluate(args, model, processor)
if args.do_predict:
predict(args, model, processor)
train_dataset = dataset.NERdataset(train_sentences, train_tags)
val_dataset = dataset.NERdataset(val_sentences, val_tags)
train_dataloader = torch.utils.data.DataLoader(
dataset = train_dataset,
batch_size = config.TRAIN_BATCH_SIZE
)
val_dataloader = torch.utils.data.DataLoader(
dataset = val_dataset,
batch_size = config.VALID_BATCH_SIZE
)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = NERModel(num_tags)
model.to(device)
param_optimizer = list(model.named_parameters())
no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
optimizer_parameters = [
{
"params": [
p for n, p in param_optimizer if not any(nd in n for nd in no_decay)
],
"weight_decay": 0.001,
},
{
"params": [
p for n, p in param_optimizer if any(nd in n for nd in no_decay)
],
"weight_decay": 0.0,