(total_loss, logits, trans,
pred_ids) = create_model(bert_config=bert_config,
is_training=False,
input_ids=input_ids_p,
input_mask=input_mask_p,
segment_ids=None,
labels=None,
num_labels=num_labels,
use_one_hot_embeddings=False,
dropout_rate=1.0)
saver = tf.train.Saver()
saver.restore(sess, tf.train.latest_checkpoint(model_dir))
tokenizer = tokenization.FullTokenizer(vocab_file=os.path.join(
bert_dir, 'vocab.txt'),
do_lower_case=args.do_lower_case)
def predict_online():
"""
do online prediction. each time make prediction for one instance.
you can change to a batch if you want.
:param line: a list. element is: [dummy_label,text_a,text_b]
:return:
"""
def convert(line):
feature = convert_single_example(0, line, label_list,
args.max_seq_length, tokenizer, 'p')
input_ids = np.reshape([feature.input_ids],
def main(_):
tf.logging.set_verbosity(tf.logging.INFO)
layer_indexes = [int(x) for x in FLAGS.layers.split(",")]
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
is_per_host = tf.contrib.tpu.InputPipelineConfig.PER_HOST_V2
run_config = tf.contrib.tpu.RunConfig(
master=FLAGS.master,
tpu_config=tf.contrib.tpu.TPUConfig(
num_shards=FLAGS.num_tpu_cores,
per_host_input_for_training=is_per_host))
examples = read_examples(FLAGS.input_file)
features = convert_examples_to_features(examples=examples,
seq_length=FLAGS.max_seq_length,
tokenizer=tokenizer)
unique_id_to_feature = {}
for feature in features:
unique_id_to_feature[feature.unique_id] = feature
model_fn = model_fn_builder(
bert_config=bert_config,
init_checkpoint=FLAGS.init_checkpoint,
layer_indexes=layer_indexes,
use_tpu=FLAGS.use_tpu,
use_one_hot_embeddings=FLAGS.use_one_hot_embeddings)
# If TPU is not available, this will fall back to normal Estimator on CPU
# or GPU.
estimator = tf.contrib.tpu.TPUEstimator(
use_tpu=FLAGS.use_tpu,
model_fn=model_fn,
config=run_config,
predict_batch_size=FLAGS.batch_size)
input_fn = input_fn_builder(features=features,
seq_length=FLAGS.max_seq_length)
with codecs.getwriter("utf-8")(tf.gfile.Open(FLAGS.output_file,
"w")) as writer:
for result in estimator.predict(input_fn, yield_single_examples=True):
unique_id = int(result["unique_id"])
feature = unique_id_to_feature[unique_id]
output_json = collections.OrderedDict()
output_json["linex_index"] = unique_id
all_features = []
for (i, token) in enumerate(feature.tokens):
all_layers = []
for (j, layer_index) in enumerate(layer_indexes):
layer_output = result["layer_output_%d" % j]
layers = collections.OrderedDict()
layers["index"] = layer_index
layers["values"] = [
round(float(x), 6)
for x in layer_output[i:(i + 1)].flat
]
all_layers.append(layers)
features = collections.OrderedDict()
features["token"] = token
features["layers"] = all_layers
all_features.append(features)
output_json["features"] = all_features
writer.write(json.dumps(output_json) + "\n")
def main(_):
if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_predict:
raise ValueError(
"At least one of `do_train`, `do_eval` or `do_predict' must be True."
)
bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
if FLAGS.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(FLAGS.max_seq_length, bert_config.max_position_embeddings))
if not os.path.exists(FLAGS.output_dir):
os.makedirs(FLAGS.output_dir)
processor = ThuProcessor()
#定义分词器
tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file,
do_lower_case=FLAGS.do_lower_case)
# estimator 运行参数
run_config = tf.estimator.RunConfig(
model_dir=FLAGS.output_dir,
save_summary_steps=FLAGS.save_summary_steps,
save_checkpoints_steps=FLAGS.save_checkpoints_steps,
keep_checkpoint_max=5,
log_step_count_steps=500,
session_config=tf.ConfigProto(log_device_placement=True)
#session_config=tf.ConfigProto(log_device_placement=True,
# device_count={'GPU': 1}))
)
train_examples = None
num_train_steps = None
num_warmup_steps = None
if FLAGS.do_train:
train_examples = processor.get_train_examples(FLAGS.data_dir)
num_train_steps = int(
len(train_examples) / FLAGS.train_batch_size *
FLAGS.num_train_epochs)
num_warmup_steps = int(num_train_steps * FLAGS.warmup_proportion)
# get_labels() must be called after get_train_examoles or other examples
label_list = processor.get_labels()
logger.info('************ label_list=', ' '.join(label_list))
model_fn = model_fn_builder(bert_config=bert_config,
num_labels=len(label_list),
init_checkpoint=FLAGS.init_checkpoint,
learning_rate=FLAGS.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps)
# params是一个dict 里面的key是model_fn 里面用到的参数名称，value是对应的数据
params = {
'batch_size': FLAGS.train_batch_size,
}
estimator = tf.estimator.Estimator(
model_fn=model_fn,
config=run_config,
params=params,
)
if FLAGS.do_train:
train_file = os.path.join(FLAGS.output_dir, "train.tf_record")
file_based_convert_examples_to_features(train_examples, label_list,
FLAGS.max_seq_length,
tokenizer, train_file, 'train')
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_examples))
logger.info(" Batch size = %d", FLAGS.train_batch_size)
logger.info(" Num steps = %d", num_train_steps)
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=FLAGS.max_seq_length,
num_label=len(label_list),
is_training=True,
drop_remainder=True)
estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
if FLAGS.do_eval:
eval_examples = processor.get_dev_examples(FLAGS.data_dir)
eval_file = os.path.join(FLAGS.output_dir, "eval.tf_record")
file_based_convert_examples_to_features(eval_examples, label_list,
FLAGS.max_seq_length,
tokenizer, eval_file, 'eval')
logger.info("***** Running evaluation *****")
logger.info(" Num examples = %d", len(eval_examples))
logger.info(" Batch size = %d", FLAGS.eval_batch_size)
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=FLAGS.max_seq_length,
num_label=len(label_list),
is_training=False,
drop_remainder=False)
result = estimator.evaluate(input_fn=eval_input_fn)
output_eval_file = os.path.join(FLAGS.output_dir, "eval_results.txt")
with tf.gfile.GFile(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])))
if FLAGS.do_predict:
predict_examples = processor.get_test_examples(FLAGS.data_dir)
predict_file = os.path.join(FLAGS.output_dir, "predict.tf_record")
file_based_convert_examples_to_features(predict_examples, label_list,
FLAGS.max_seq_length,
tokenizer, predict_file,
'test')
logger.info("***** Running prediction*****")
logger.info(" Num examples = %d", len(predict_examples))
logger.info(" Batch size = %d", FLAGS.predict_batch_size)
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=FLAGS.max_seq_length,
num_label=len(label_list),
is_training=False,
drop_remainder=False)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(FLAGS.output_dir,
"test_results.txt")
with tf.gfile.GFile(output_predict_file, "w") as writer:
logger.info("***** Predict results *****")
for prediction in result:
output_line = "\t".join(
str(class_probability)
for class_probability in prediction) + "\n"
writer.write(output_line)
def train(args):
os.environ['CUDA_VISIBLE_DEVICES'] = args.device_map
set_os_variables(args.device_map, args.verbose)
processors = {"ner": NerProcessor, "class": ClassProcessor}
bert_config = modeling.BertConfig.from_json_file(args.bert_config_file)
if args.max_seq_length > bert_config.max_position_embeddings:
raise ValueError(
"Cannot use sequence length %d because the BERT model "
"was only trained up to sequence length %d" %
(args.max_seq_length, bert_config.max_position_embeddings))
# 在re train 的时候，才删除上一轮产出的文件，在predicted 的时候不做clean
# Clean files produced in the previous round retrain, no cleaning when predicting
if args.clean and args.do_train:
if os.path.exists(args.output_dir):
def del_file(path):
ls = os.listdir(path)
for i in ls:
c_path = os.path.join(path, i)
if os.path.isdir(c_path):
del_file(c_path)
else:
os.remove(c_path)
try:
del_file(args.output_dir)
except Exception as e:
print(e)
print('pleace remove the files of output dir and data.conf')
exit(-1)
#check output dir exists
if not os.path.exists(args.output_dir):
os.mkdir(args.output_dir)
processor = processors[args.ner](args.output_dir)
tokenizer = tokenization.FullTokenizer(vocab_file=args.vocab_file,
do_lower_case=args.do_lower_case)
session_config = tf.ConfigProto(log_device_placement=False,
inter_op_parallelism_threads=0,
intra_op_parallelism_threads=0,
allow_soft_placement=True)
run_config = tf.estimator.RunConfig(model_dir=args.output_dir,
save_summary_steps=500,
save_checkpoints_steps=500,
session_config=session_config)
train_examples = None
eval_examples = None
num_train_steps = None
num_warmup_steps = None
if args.do_train and args.do_eval:
# 加载训练数据
train_examples = processor.get_train_examples(args.data_dir)
num_train_steps = int(
len(train_examples) * 1.0 / args.batch_size *
args.num_train_epochs)
if num_train_steps < 1:
raise AttributeError('training data is so small...')
num_warmup_steps = int(num_train_steps * args.warmup_proportion)
tf.logging.info("***** Running training *****")
tf.logging.info(" Num examples = %d", len(train_examples))
tf.logging.info(" Batch size = %d", args.batch_size)
tf.logging.info(" Num steps = %d", num_train_steps)
eval_examples = processor.get_dev_examples(args.data_dir)
# 打印验证集数据信息
tf.logging.info("***** Running evaluation *****")
tf.logging.info(" Num examples = %d", len(eval_examples))
tf.logging.info(" Batch size = %d", args.batch_size)
label_list = processor.get_labels()
tf.logging.info("** LABEL_LIST: {}".format(label_list))
# 返回的model_dn 是一个函数，其定义了模型，训练，评测方法，并且使用钩子参数，加载了BERT模型的参数进行了自己模型的参数初始化过程
# The returned model_dn is a function that defines the model, training, and evaluation methods, and uses hook parameters, loads the parameters of the BERT model, and performs the parameter initialization process of its own model.
# tf 新的架构方法，通过定义model_fn 函数，定义模型，然后通过EstimatorAPI进行模型的其他工作，Es就可以控制模型的训练，预测，评估工作等。
# tf new architecture method. By defining the model_fn function, defining the model, and then performing other work on the model through the EstimatorAPI, Es can control the training, prediction, and evaluation of the model.
model_fn = model_fn_builder(
bert_config=bert_config,
num_labels=len(label_list), # + 1,
init_checkpoint=args.init_checkpoint,
learning_rate=args.learning_rate,
num_train_steps=num_train_steps,
num_warmup_steps=num_warmup_steps,
args=args)
params = {'batch_size': args.batch_size}
estimator = tf.estimator.Estimator(model_fn,
params=params,
config=run_config)
if args.do_train and args.do_eval:
# 1. 将数据转化为tf_record 数据 - (en) 1. Converting data to tf_record data
train_file = os.path.join(args.output_dir, "train.tf_record")
if not os.path.exists(train_file):
_, _ = filed_based_convert_examples_to_features(
train_examples, label_list, args.max_seq_length, tokenizer,
train_file, args.output_dir)
# 2.读取record 数据，组成batch - (en) 2. Read record data to form a batch
train_input_fn = file_based_input_fn_builder(
input_file=train_file,
seq_length=args.max_seq_length,
is_training=True,
drop_remainder=True)
# estimator.train(input_fn=train_input_fn, max_steps=num_train_steps)
eval_file = os.path.join(args.output_dir, "eval.tf_record")
if not os.path.exists(eval_file):
batch_tokens, batch_labels = filed_based_convert_examples_to_features(
eval_examples, label_list, args.max_seq_length, tokenizer,
eval_file, args.output_dir)
eval_input_fn = file_based_input_fn_builder(
input_file=eval_file,
seq_length=args.max_seq_length,
is_training=False,
drop_remainder=False)
# train and eval together
# early stop hook
early_stopping_hook = tf.estimator.experimental.stop_if_no_decrease_hook(
estimator=estimator,
metric_name='loss',
max_steps_without_decrease=num_train_steps,
eval_dir=None,
min_steps=0,
run_every_secs=None,
run_every_steps=args.save_checkpoints_steps)
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn,
max_steps=num_train_steps,
hooks=[early_stopping_hook])
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
if args.do_predict:
token_path = os.path.join(args.output_dir, "token_test.txt")
if os.path.exists(token_path):
os.remove(token_path)
with codecs.open(os.path.join(args.output_dir, 'label2id.pkl'),
'rb') as rf:
label2id = pickle.load(rf)
id2label = {value: key for key, value in label2id.items()}
tf.logging.info("****** id2label ****** \n {}\n".format(id2label))
predict_examples = processor.get_test_examples(args.data_dir)
predict_file = os.path.join(args.output_dir, "predict.tf_record")
batch_tokens, batch_labels = filed_based_convert_examples_to_features(
predict_examples,
label_list,
args.max_seq_length,
tokenizer,
predict_file,
args.output_dir,
mode="test")
tf.logging.info("***** Running prediction*****")
tf.logging.info(" Num examples = %d", len(predict_examples))
tf.logging.info(" Batch size = %d", args.batch_size)
predict_drop_remainder = False
predict_input_fn = file_based_input_fn_builder(
input_file=predict_file,
seq_length=args.max_seq_length,
is_training=False,
drop_remainder=predict_drop_remainder)
result = estimator.predict(input_fn=predict_input_fn)
output_predict_file = os.path.join(args.output_dir, "label_test.txt")
Writer(output_predict_file, result, batch_tokens, batch_labels,
id2label)
"""
def result_to_pair(writer):
for predict_line, prediction in zip(predict_examples, result):
idx = 0
line = ''
line_token = str(predict_line.text).split(' ')
label_token = str(predict_line.label).split(' ')
len_seq = len(label_token)
if len(line_token) != len(label_token):
tf.logging.info(predict_line.text)
tf.logging.info(predict_line.label)
break
for id in prediction:
if idx >= len_seq:
break
if id == 0:
continue
curr_labels = id2label[id]
if curr_labels in ['[CLS]', '[SEP]']:
continue
try:
line += line_token[idx] + ' ' + label_token[idx] + ' ' + curr_labels + '\n'
except Exception as e:
tf.logging.info(e)
tf.logging.info(predict_line.text)
tf.logging.info(predict_line.label)
line = ''
break
idx += 1
writer.write(line + '\n')
with codecs.open(output_predict_file, 'w', encoding='utf-8') as writer:
result_to_pair(writer)
"""
from bert_as_server.train import conlleval
eval_result = conlleval.return_report(output_predict_file)
print(''.join(eval_result))
# 写结果到文件中
with codecs.open(os.path.join(args.output_dir, 'predict_score.txt'),
'a',
encoding='utf-8') as fd:
fd.write(''.join(eval_result))
# filter model
if args.filter_adam_var:
adam_filter(args.output_dir)