(args_dict)
| 80 | relu_dropout=0.1) |
| 81 | |
| 82 | def translate(args_dict): |
| 83 | |
| 84 | batch_size = args_dict['batch_size'] |
| 85 | beam_size = args_dict['beam_width'] |
| 86 | max_seq_len = args_dict['max_seq_len'] |
| 87 | model_dir = args_dict["model_dir"] |
| 88 | source_file = args_dict["source"] |
| 89 | tgt_file = args_dict["target"] |
| 90 | time_args = args_dict["test_time"] |
| 91 | beam_search_diversity_rate = args_dict['beam_search_diversity_rate'] |
| 92 | sampling_topk = args_dict['sampling_topk'] |
| 93 | sampling_topp = args_dict['sampling_topp'] |
| 94 | tf_datatype = tf.float32 |
| 95 | max_ite = args_dict['max_iteration'] |
| 96 | if args_dict['data_type'] == "fp16": |
| 97 | tf_datatype = tf.float16 |
| 98 | elif args_dict['data_type'] == "bf16": |
| 99 | tf_datatype = tf.bfloat16 |
| 100 | |
| 101 | print("\n=============== Argument ===============") |
| 102 | for key in args_dict: |
| 103 | print("{}: {}".format(key, args_dict[key])) |
| 104 | print("========================================") |
| 105 | |
| 106 | # Define the "base" Transformer model. |
| 107 | source_inputter = onmt.inputters.WordEmbedder("source_vocabulary", embedding_size=512, dtype=tf_datatype) |
| 108 | target_inputter = onmt.inputters.WordEmbedder("target_vocabulary", embedding_size=512, dtype=tf_datatype) |
| 109 | |
| 110 | inputter = onmt.inputters.ExampleInputter(source_inputter, target_inputter) |
| 111 | inputter.initialize({ |
| 112 | "source_vocabulary": args_dict["source_vocabulary"], |
| 113 | "target_vocabulary": args_dict["target_vocabulary"] |
| 114 | }) |
| 115 | |
| 116 | mode = tf.estimator.ModeKeys.PREDICT |
| 117 | |
| 118 | np.random.seed(1) |
| 119 | tf.set_random_seed(1) |
| 120 | |
| 121 | # Create the inference dataset. |
| 122 | dataset = inputter.make_inference_dataset(source_file, batch_size) |
| 123 | iterator = dataset.make_initializable_iterator() |
| 124 | source = iterator.get_next() |
| 125 | |
| 126 | encoder_args = TransformerArgument(beam_width=1, |
| 127 | head_num=NUM_HEADS, |
| 128 | size_per_head=SIZE_PER_HEAD, |
| 129 | inter_size=NUM_HEADS*SIZE_PER_HEAD*4, |
| 130 | num_layer=NUM_LAYERS, |
| 131 | dtype=tf_datatype, |
| 132 | remove_padding=True, |
| 133 | allow_gemm_test=False) |
| 134 | |
| 135 | # Encode the source. |
| 136 | with tf.variable_scope("transformer/encoder"): |
| 137 | source_embedding = source_inputter.make_inputs(source) |
| 138 | source_embedding = tf.cast(source_embedding, tf_datatype) |
| 139 |
no test coverage detected