def __init__(self, config):
super(TextCNN, self).__init__()
self.name_to_features = {
"input": tf.FixedLenFeature([config['data']['n_steps']], tf.int64),
"target": tf.FixedLenFeature([1], tf.int64)
}
self.shuffle_size = config['data']['shuffle_size']
self.num_parallel_calls = config['data']['num_parallel_calls']
self.batch_size = config['data']['batch_size']
self.train_drop_remainder = config['data']['train_drop_remainder']
self.prefetch_size = config['data']['prefetch_size']
self.vocab = Vocab.load(config['model']['vocab'])
self.emb_size = config['model']['emb_size']
self.filter_sizes = config['model']['filter_sizes']
self.n_filters = config['model']['n_filters']
self.dropout = config['model']['dropout']
self.n_classes = config['model']['n_classes']
def __init__(self, config):
super(TextCNN, self).__init__()
self.name_to_features = {
"input": tf.FixedLenFeature([config['data']['n_steps']], tf.int64),
"target": tf.FixedLenFeature([1], tf.int64)
}
self.shuffle_size = config['data']['shuffle_size']
self.num_parallel_calls = config['data']['num_parallel_calls']
self.batch_size = config['data']['batch_size']
self.train_drop_remainder = config['data']['train_drop_remainder']
self.prefetch_size = config['data']['prefetch_size']
self.vocab = Vocab.load(config['model']['vocab'])
self.emb_size = config['model']['emb_size']
self.rnn_units = config['model']['rnn_units']
self.rnn_layers = config['model']['rnn_layers']
self.rnn_bidi = config['model']['rnn_bidi']
# if rnn_residual is true, emb_size should be equal to rnn_units
self.rnn_residual = config['model']['rnn_residual']
self.rnn_state_scheme = config['model']['rnn_state_scheme']
self.dropout = config['model']['dropout']
self.n_classes = config['model']['n_classes']
def main(model_dir, ckpt, device, emb_path):
model_dir = Path(model_dir)
config = Config(model_dir / 'config.json')
with open(config.vocab.words) as fwords, open(
config.vocab.labels) as flabels:
vocab = Vocab.load(fwords, flabels)
model = MultiClassModel(vocab, config.model)
logger.info("Loading parameters")
model.load_state_dict(torch.load(model_dir / ckpt))
if emb_path:
logger.info("Loading embeddings")
with open(emb_path) as femb, open(config.vocab.labels) as flabels:
vocab, vecs = Vocab.build_from_emb(femb, flabels)
model.set_embedding(vocab, vecs)
if device is not None:
device = torch.device(device)
else:
device = torch.device(config.training.device)
model = model.to(device)
model.eval()
logger.info("Start prediction")
input_lines = []
batch_tokens = []
for i, line in enumerate(sys.stdin):
input_lines.append(line)
tokens = line.split()
tokens = vocab.idfy(tokens)
batch_tokens.append(tokens)
if (i + 1) % 32 == 0:
batch_padded_tokens, batch_mask = batch_pad(batch_tokens, vocab)
batch_padded_tokens = torch.tensor(batch_padded_tokens,
dtype=torch.long,
requires_grad=False).to(device)
batch_mask = torch.tensor(batch_mask,
requires_grad=False).to(device)
result = model(batch_padded_tokens, batch_mask, predict=True)
prob = torch.nn.functional.softmax(result['logits'], dim=1)
for i in range(len(result['pred'])):
print(
json.dumps({
'input':
input_lines[i],
'prob': {
vocab.id2label(j): float(prob[i, j])
for j in range(prob.shape[1])
},
'prediction':
vocab.id2label(int(result['pred'][i]))
}))
input_lines = []
batch_tokens = []
if not input_lines:
return
batch_padded_tokens, batch_mask = batch_pad(batch_tokens, vocab)
result = model(batch_padded_tokens, batch_mask, predict=True)
prob = torch.nn.functional.softmax(result['logits'], dim=1)
for i in len(result['pred']):
print(
json.dumps({
'input': input_lines[i],
'prob': {
vocab.id2label(j): float(prob[i, j])
for j in range(prob.shape[1])
},
'prediction': vocab.id2label(int(result['pred'][i]))
}))
def main(args):
print(args)
with codecs.open(args.config, 'r', "utf8") as fin:
config = yaml.load(fin, Loader=yaml.SafeLoader)
model = TextCNN(config)
session_config = tf.ConfigProto(
allow_soft_placement=True,
gpu_options=tf.GPUOptions(allow_growth=True)
)
run_config = tf.estimator.RunConfig(
session_config=session_config,
tf_random_seed=config['train']['tf_random_seed'],
save_checkpoints_steps=config['train']['save_checkpoints_steps'],
keep_checkpoint_max=config['train']['keep_checkpoint_max'],
save_summary_steps=config['train']['save_summary_steps']
)
estimator = tf.estimator.Estimator(
model_fn=model.model_fn,
model_dir=config['train']['model_dir'],
config=run_config,
params={}
)
for epoch in range(config['train']['epochs']):
if args.train:
estimator.train(
input_fn=model.get_train_input_fn(config['train']['train_dir']),
steps=config['train']['steps_per_train']
)
if args.eval:
res = estimator.evaluate(
input_fn=model.get_eval_input_fn(config['train']['eval_dir']),
steps=config['train']['steps_per_eval']
)
dir_out = os.path.dirname(config['train']['eval_result_dir'])
if not os.path.exists(dir_out):
os.makedirs(dir_out)
with codecs.open(config['train']['eval_result_dir'], 'a', 'utf8') as fout:
tp, tn, fp, fn = res['tp'], res['tn'], res['fp'], res['fn']
accuracy = (tp + tn) / (tp + tn + fp + fn)
precision = tp / (tp + fp)
recall = tp / (tp + fn)
f1 = (2 * precision * recall) / (precision + recall)
res_str = 'global_step: %d\n' \
'loss: %.7f\n' \
'TP: %d TN: %d FP: %d FN: %d\n' \
'accuracy: %.6f\n' \
'precision: %.6f\n' \
'recall: %.6f\n' \
'f1: %.6f\n\n\n' \
% (res['global_step'], res['loss'], tp, tn, fp, fn, accuracy, precision, recall, f1)
print(res_str)
fout.write(res_str)
if args.test:
vocab = Vocab.load(config['model']['vocab'])
dir_out = os.path.dirname(config['train']['test_result_dir'])
if not os.path.exists(dir_out):
os.makedirs(dir_out)
with codecs.open(config['train']['test_result_dir'], 'w', 'utf8') as fout:
for pred in estimator.predict(
input_fn=model.get_infer_input_fn(config['train']['test_dir']),
checkpoint_path=args.ckpt
):
utt = pred_2_str(pred, vocab)
fout.write(utt + '\n')
