def _process_type_midd(condition, decoder, grammar_stack, next_inputs, predicted_ids): """Process when output type is MID Args: condition (TYPE): NULL decoder (TYPE): NULL grammar_stack (TYPE): NULL next_inputs (TYPE): NULL predicted_ids (TYPE): NULL Returns: TODO Raises: NULL """ midd_pred_ids = fluider.elementwise_mul(predicted_ids, condition, axis=0, force=True) ## get grammar desc # 解码结果(语法ID)对应的具体语法规则序列。比如解码结果为 SingleSQL,则对应的语法序列为 Select Filter # shape = [batch_size, grammar.max_desc_len] gmr_desc = decoder.grammar_desc(midd_pred_ids) # 语法规则序列的长度,比如 SingleSQL --> Select Filter, 则长度为2 # shape = [batch_size, 1] gmr_desc_lens = decoder.grammar_desc_lens(midd_pred_ids) # shape = [batch_size, 1] gmr_desc_pos = tensor.zeros_like(gmr_desc_lens) ## generate next grammar mask by first token in desc next_output = nn_utils.batch_gather(gmr_desc, gmr_desc_pos) next_actions = decoder.grammar_action(next_output) next_gmr_mask = decoder.grammar_mask(next_output) ## push left grammar tokens to stack gmr_stack_tmp, gmr_stack_pos_tmp = _push_to_stack(gmr_desc, gmr_desc_pos, gmr_desc_lens, grammar_stack) ## save result, while condition is True new_gmr_stack, new_gmr_stack_pos, new_actions, new_gmr_mask = nn_utils.ifelse( condition, [gmr_stack_tmp, gmr_stack_pos_tmp, next_actions, next_gmr_mask], [ grammar_stack.data, grammar_stack.pos, next_inputs.action, next_inputs.gmr_mask ]) layers.utils.map_structure( layers.assign, [new_gmr_stack, new_gmr_stack_pos, new_actions, new_gmr_mask], [ grammar_stack.data, grammar_stack.pos, next_inputs.action, next_inputs.gmr_mask ])
def _process_type_leaf(condition, decoder, grammar_stack, next_inputs, finished): """Process when output type is LEAF Args: condition (TYPE): NULL decoder (TYPE): NULL grammar_stack (StackData): (gmr_stack_data, gmr_stack_pos) next_inputs (DecoderInputsWrapper): (input_var, action, grammar_mask) finished (TYPE): NULL Returns: None Raises: NULL """ ## pop stack next_output, valid_pos, gmr_stack_tmp = data_structure.Stack.pop( grammar_stack, mask=True, in_place=False) valid_pos = fluider.squeeze(valid_pos, [1]) ## update next grammar mask next_actions = layers.elementwise_mul(decoder.grammar_action(next_output), layers.cast( valid_pos, dtype=next_inputs.action.dtype), axis=0) next_gmr_mask = layers.elementwise_mul( decoder.grammar_mask(next_output), layers.cast(valid_pos, dtype=next_inputs.gmr_mask.dtype), axis=0) ## save result, while condition is True new_gmr_stack_data, new_gmr_stack_pos, new_actions, new_gmr_mask = nn_utils.ifelse( condition, [gmr_stack_tmp.data, gmr_stack_tmp.pos, next_actions, next_gmr_mask], [ grammar_stack.data, grammar_stack.pos, next_inputs.action, next_inputs.gmr_mask ]) layers.utils.map_structure( layers.assign, [new_gmr_stack_data, new_gmr_stack_pos, next_actions, new_gmr_mask], [ grammar_stack.data, grammar_stack.pos, next_inputs.action, next_inputs.gmr_mask ]) layers.logical_or(finished, layers.logical_and(condition, layers.logical_not(valid_pos)), out=finished)
def _save_predict_output(outputs_array, predicted_ids, finished): """save predicted_ids to outputs_array, while finished is false. Args: outputs_array (TYPE): NULL predicted_ids (TYPE): NULL finished (TYPE): NULL Returns: TODO Raises: NULL """ out_arr_tmp = data_structure.Array.push(outputs_array, predicted_ids, in_place=False) new_arr_data, new_arr_pos = nn_utils.ifelse( finished, [outputs_array.data, outputs_array.pos], [out_arr_tmp.data, out_arr_tmp.pos]) layers.utils.map_structure(layers.assign, [new_arr_data, new_arr_pos], [outputs_array.data, outputs_array.pos])
def _push_to_stack(gmr_desc, gmr_pos, gmr_lens, gmr_stack_info): """push grammar id in gmr_desc from gmr_pos to gmr_lens to gmr_stack. and update step_gmr_pos Args: gmr_desc (TYPE): NULL gmr_pos (TYPE): NULL gmr_lens (TYPE): NULL gmr_stack_info (tuple): [in/out] (gmr_stack, gmr_stack_pos) Returns: tuple (gmr_stack, gmr_stack_pos) Raises: NULL """ gmr_stack, gmr_stack_pos = gmr_stack_info mv_step = layers.cast(layers.greater_than(gmr_lens, layers.zeros_like(gmr_lens)), dtype=gmr_lens.dtype) gmr_mv_pos = layers.elementwise_sub(gmr_lens, mv_step) cond = layers.reduce_any(layers.greater_than(gmr_mv_pos, gmr_pos)) while_op = layers.While(cond) with while_op.block(): gmr_ids = nn_utils.batch_gather(gmr_desc, gmr_mv_pos) gmr_stack_tmp, gmr_stack_pos_tmp = data_structure.Stack.push( gmr_stack_info, gmr_ids, in_place=False) mv_cond = layers.greater_than(gmr_mv_pos, gmr_pos) gmr_mv_pos_tmp = fluider.elementwise_sub(gmr_mv_pos, mv_cond, force=True) new_gmr_stack, new_gmr_stack_pos = nn_utils.ifelse( mv_cond, [gmr_stack_tmp, gmr_stack_pos_tmp], [gmr_stack, gmr_stack_pos]) layers.utils.map_structure(layers.assign, [new_gmr_stack, new_gmr_stack_pos], [gmr_stack, gmr_stack_pos]) layers.assign(gmr_mv_pos_tmp, gmr_mv_pos) layers.assign( layers.reduce_any(layers.greater_than(gmr_mv_pos, gmr_pos)), cond) return gmr_stack, gmr_stack_pos
def _check_finished(decoder, next_inputs, finished, outputs_array): """check finished instance by next_inputs.action, and update finished tag and write END to outputs Args: decoder (TYPE): NULL next_inputs (TYPE): NULL finished (TYPE): NULL outputs_array (TYPE): NULL Returns: TODO Raises: NULL """ act_stop = tensor.fill_constant_batch_size_like( next_inputs.action, shape=next_inputs.action.shape, value=decoder._grammar.ACTION_STOP, dtype='int64') new_finished = layers.logical_and( layers.equal(next_inputs.action, act_stop), layers.logical_not(finished)) end_token_id = tensor.fill_constant_batch_size_like( outputs_array.data, shape=[-1], value=decoder._grammar.END, dtype=outputs_array.data.dtype) out_data_tmp, out_pos_tmp = data_structure.Array.push(outputs_array, end_token_id, in_place=False) new_data, new_pos = nn_utils.ifelse( new_finished, [out_data_tmp, out_pos_tmp], [outputs_array.data, outputs_array.pos]) layers.assign(new_data, outputs_array.data) layers.assign(new_pos, outputs_array.pos) layers.logical_or(finished, new_finished, out=finished)