def main(_): index = FLAGS.index head = FLAGS.head vocab_path = FLAGS.vocab_path attention_file = FLAGS.attention_file subtokenizer = tokenization.restore_subtokenizer_from_vocab_files(vocab_path) subtoken_list = subtokenizer._subtoken_list attn = np.load(attention_file, allow_pickle=True).item() src_src = attn['src_src_attention'] tgt_src = attn['tgt_src_attention'] tgt_tgt = attn['tgt_tgt_attention'] src = attn['src'] tgt = attn['tgt'] draw_attention_weights( src[index], src[index], src_src[index, head], subtoken_list) plt.savefig('src_src.png', dpi=256) print('src_src_attention saved to "src_src.png".') draw_attention_weights( tgt[index], src[index], tgt_src[index, head], subtoken_list) plt.savefig('tgt_src.png', dpi=256) print('tgt_src_attention saved to "tgt_src.png".') draw_attention_weights( tgt[index], tgt[index], tgt_tgt[index, head], subtoken_list) plt.savefig('tgt_tgt.png', dpi=256) print('tgt_tgt_attention saved to "tgt_tgt.png".')
def main(_): vocab_path = FLAGS.vocab_path model_dir = FLAGS.model_dir encoder_stack_size = FLAGS.encoder_stack_size decoder_stack_size = FLAGS.decoder_stack_size hidden_size = FLAGS.hidden_size num_heads = FLAGS.num_heads filter_size = FLAGS.filter_size dropout_rate = FLAGS.dropout_rate extra_decode_length = FLAGS.extra_decode_length beam_width = FLAGS.beam_width alpha = FLAGS.alpha decode_batch_size = FLAGS.decode_batch_size src_max_length = FLAGS.src_max_length source_text_filename = FLAGS.source_text_filename target_text_filename = FLAGS.target_text_filename translation_output_filename = FLAGS.translation_output_filename # transformer model subtokenizer = tokenization.restore_subtokenizer_from_vocab_files( vocab_path) vocab_size = subtokenizer.vocab_size model = TransformerModel(vocab_size=vocab_size, encoder_stack_size=encoder_stack_size, decoder_stack_size=decoder_stack_size, hidden_size=hidden_size, num_heads=num_heads, filter_size=filter_size, dropout_rate=dropout_rate, extra_decode_length=extra_decode_length, beam_width=beam_width, alpha=alpha) ckpt = tf.train.Checkpoint(model=model) latest_ckpt = tf.train.latest_checkpoint(model_dir) if latest_ckpt is None: raise ValueError('No checkpoint is found in %s' % model_dir) print('Loaded latest checkpoint ', latest_ckpt) ckpt.restore(latest_ckpt).expect_partial() # build evaluator evaluator = SequenceTransducerEvaluator(model, subtokenizer, decode_batch_size, src_max_length) # translates input sequences, and optionally evaluates BLEU score if # groundtruth target sequences are provided if target_text_filename is not None: case_insensitive_score, case_sensitive_score = evaluator.evaluate( source_text_filename, target_text_filename, translation_output_filename) print('BLEU(case insensitive): %f' % case_insensitive_score) print('BLEU(case sensitive): %f' % case_sensitive_score) else: evaluator.translate(source_text_filename, translation_output_filename) print( 'Inference mode: no groundtruth translations.\nTranslations written ' 'to file "%s"' % translation_output_filename)
def main(_): filename = FLAGS.filename vocab_path = FLAGS.vocab_path model_dir = FLAGS.model_dir m_seq_len = FLAGS.m_seq_len cutoffs = FLAGS.cutoffs adaptive_embedding = FLAGS.adaptive_embedding stack_size = FLAGS.stack_size hidden_size = FLAGS.hidden_size num_heads = FLAGS.num_heads filter_size = FLAGS.filter_size tie_biases = FLAGS.tie_biases with tf.io.gfile.GFile(filename + '.json') as f: dataset_config = json.load(f) subword = dataset_config['subword'] batch_size = dataset_config['batch_size'] if subword: tokenizer = tokenization.restore_subtokenizer_from_vocab_files(vocab_path) else: tokenizer = tokenization.restore_tokenizer_from_vocab_files(vocab_path) vocab_size = tokenizer.vocab_size cutoffs = list(map(int, cutoffs)) model = TransformerXLModel(adaptive_embedding, vocab_size, cutoffs, stack_size, hidden_size, num_heads, filter_size, tie_biases=tie_biases) parse_fn = functools.partial(parse_fn_sequence_pair, keys=('inputs', 'labels'), dtype='int32') dataset = tf.data.TFRecordDataset(filename + '.tfrecord') dataset = dataset.map(parse_fn).batch(batch_size) ckpt = tf.train.Checkpoint(model=model) latest_ckpt = tf.train.latest_checkpoint(model_dir) if latest_ckpt is None: raise ValueError('No checkpoint is found in %s' % model_dir) print('Loaded latest checkpoint', latest_ckpt) ckpt.restore(latest_ckpt).expect_partial() evaluator = TransformerXLModelEvaluator(model, m_seq_len, batch_size, vocab_size, adaptive_embedding) print('Evaluating file %s...' % (filename + '.tfrecord')) perplexity = evaluator.evaluate(dataset) print('Perplexity:', perplexity)
def main(_): data_dir = FLAGS.data_dir vocab_path = FLAGS.vocab_path model_dir = FLAGS.model_dir encoder_stack_size = FLAGS.encoder_stack_size decoder_stack_size = FLAGS.decoder_stack_size hidden_size = FLAGS.hidden_size num_heads = FLAGS.num_heads filter_size = FLAGS.filter_size dropout_rate = FLAGS.dropout_rate max_num_tokens = FLAGS.max_num_tokens max_length = FLAGS.max_length num_parallel_calls = FLAGS.num_parallel_calls learning_rate = FLAGS.learning_rate learning_rate_warmup_steps = FLAGS.learning_rate_warmup_steps optimizer_adam_beta1 = FLAGS.optimizer_adam_beta1 optimizer_adam_beta2 = FLAGS.optimizer_adam_beta2 optimizer_adam_epsilon = FLAGS.optimizer_adam_epsilon label_smoothing = FLAGS.label_smoothing num_steps = FLAGS.num_steps save_ckpt_per_steps = FLAGS.save_ckpt_per_steps # transformer model subtokenizer = tokenization.restore_subtokenizer_from_vocab_files( vocab_path) vocab_size = subtokenizer.vocab_size model = TransformerModel(vocab_size=vocab_size, encoder_stack_size=encoder_stack_size, decoder_stack_size=decoder_stack_size, hidden_size=hidden_size, num_heads=num_heads, filter_size=filter_size, dropout_rate=dropout_rate) # training dataset builder = dataset.DynamicBatchDatasetBuilder(max_num_tokens, True, max_length, num_parallel_calls) filenames = sorted(glob.glob(os.path.join(data_dir, SUFFIX))) train_ds = builder.build_dataset(filenames) # learning rate and optimizer optimizer = tf.keras.optimizers.Adam(utils.LearningRateSchedule( learning_rate, hidden_size, learning_rate_warmup_steps), optimizer_adam_beta1, optimizer_adam_beta2, epsilon=optimizer_adam_epsilon) # checkpoint ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer) # build trainer and start training trainer = SequenceTransducerTrainer(model, label_smoothing) trainer.train(train_ds, optimizer, ckpt, model_dir, num_steps, save_ckpt_per_steps)
def main(_): filename = FLAGS.filename vocab_path = FLAGS.vocab_path model_dir = FLAGS.model_dir m_seq_len = FLAGS.m_seq_len cutoffs = FLAGS.cutoffs adaptive_embedding = FLAGS.adaptive_embedding stack_size = FLAGS.stack_size hidden_size = FLAGS.hidden_size num_heads = FLAGS.num_heads filter_size = FLAGS.filter_size dropout_rate = FLAGS.dropout_rate dropout_rate_attention = FLAGS.dropout_rate_attention tie_biases = FLAGS.tie_biases learning_rate = FLAGS.learning_rate learning_rate_warmup_steps = FLAGS.learning_rate_warmup_steps optimizer_adam_beta1 = FLAGS.optimizer_adam_beta1 optimizer_adam_beta2 = FLAGS.optimizer_adam_beta2 optimizer_adam_epsilon = FLAGS.optimizer_adam_epsilon warmup_lr = FLAGS.warmup_lr clip_norm = FLAGS.clip_norm alpha = FLAGS.alpha num_steps = FLAGS.num_steps save_ckpt_per_steps = FLAGS.save_ckpt_per_steps with tf.io.gfile.GFile(filename + '.json') as f: dataset_config = json.load(f) subword = dataset_config['subword'] batch_size = dataset_config['batch_size'] # transformerxl model if subword: tokenizer = tokenization.restore_subtokenizer_from_vocab_files( vocab_path) else: tokenizer = tokenization.restore_tokenizer_from_vocab_files(vocab_path) vocab_size = tokenizer.vocab_size cutoffs = list(map(int, cutoffs)) model = TransformerXLModel(adaptive_embedding, vocab_size, cutoffs, stack_size, hidden_size, num_heads, filter_size, dropout_rate=dropout_rate, dropout_rate_attention=dropout_rate_attention, tie_biases=tie_biases) # training datset parse_fn = functools.partial(parse_fn_sequence_pair, keys=('inputs', 'labels'), dtype='int32') dataset = tf.data.TFRecordDataset(filename + '.tfrecord') dataset = dataset.map(parse_fn).repeat().batch(batch_size) # learning rate and optimizer schedule = CosineDecayLearningRateSchedule( learning_rate=learning_rate, decay_steps=num_steps - learning_rate_warmup_steps, alpha=alpha, warmup_steps=learning_rate_warmup_steps, warmup_lr=warmup_lr) optimizer = tf.keras.optimizers.Adam(schedule, optimizer_adam_beta1, optimizer_adam_beta2, epsilon=optimizer_adam_epsilon) # checkpoint ckpt = tf.train.Checkpoint(model=model, optimizer=optimizer) # build trainer and start training trainer = TransformerXLModelTrainer(model, m_seq_len, batch_size, vocab_size, adaptive_embedding) trainer.train(dataset, optimizer, ckpt, model_dir, num_steps, save_ckpt_per_steps, clip_norm)
def main(_): prompt_filename = FLAGS.prompt_filename filename = FLAGS.filename vocab_path = FLAGS.vocab_path model_dir = FLAGS.model_dir decoding_method = FLAGS.decoding_method m_seq_len = FLAGS.m_seq_len cutoffs = FLAGS.cutoffs adaptive_embedding = FLAGS.adaptive_embedding num_tokens = FLAGS.num_tokens stack_size = FLAGS.stack_size hidden_size = FLAGS.hidden_size num_heads = FLAGS.num_heads filter_size = FLAGS.filter_size tie_biases = FLAGS.tie_biases with tf.io.gfile.GFile(filename + '.json') as f: dataset_config = json.load(f) subword = dataset_config['subword'] if subword: tokenizer = tokenization.restore_subtokenizer_from_vocab_files(vocab_path) else: tokenizer = tokenization.restore_tokenizer_from_vocab_files(vocab_path) vocab_size = tokenizer.vocab_size cutoffs = list(map(int, cutoffs)) model = TransformerXLModel(adaptive_embedding, vocab_size, cutoffs, stack_size, hidden_size, num_heads, filter_size, tie_biases=tie_biases) ckpt = tf.train.Checkpoint(model=model) latest_ckpt = tf.train.latest_checkpoint(model_dir) if latest_ckpt is None: raise ValueError('No checkpoint is found in %s' % model_dir) print('Loaded latest checkpoint', latest_ckpt) ckpt.restore(latest_ckpt).expect_partial() inferencer = TransformerXLModelInferencer(model, m_seq_len, 1, vocab_size, adaptive_embedding, decoding_method, num_tokens=num_tokens) with open(prompt_filename) as f: prompt = f.read() prompt_token_ids = tokenizer.encode(prompt, add_eos=False) token_id_list = inferencer.infer(tf.constant([prompt_token_ids])) text = tokenizer.decode(token_id_list) print('\nPrompted Sequence:\n') print(prompt, '\n\n') print('Predicted sequence:\n') print(text)