def _batch_predict_entities(self, sequence_feature_lengths: tf.Tensor, text_transformed: tf.Tensor) -> Dict[(Text, tf.Tensor)]:
    predictions: Dict[(Text, tf.Tensor)] = {}
    entity_tags = None
    for tag_spec in self._entity_tag_specs:
        if (tag_spec.num_tags == 0):
            continue
        name = tag_spec.tag_name
        _input = text_transformed
        if (entity_tags is not None):
            _tags = self._tf_layers[f'embed.{name}.tags'](entity_tags)
            _input = tf.concat([_input, _tags], axis=(- 1))
        _logits = self._tf_layers[f'embed.{name}.logits'](_input)
        (pred_ids, confidences) = self._tf_layers[f'crf.{name}'](_logits, sequence_feature_lengths)
        predictions[f'e_{name}_ids'] = pred_ids
        predictions[f'e_{name}_scores'] = confidences
        if (name == ENTITY_ATTRIBUTE_TYPE):
            entity_tags = tf.one_hot(tf.cast(pred_ids, tf.int32), depth=tag_spec.num_tags)
    return predictions
def _batch_predict_entities(self, sequence_feature_lengths: tf.Tensor, text_transformed: tf.Tensor) -> Dict[(Text, tf.Tensor)]:
    predictions: Dict[(Text, tf.Tensor)] = {}
    entity_tags = None
    for tag_spec in self._entity_tag_specs:
        if (tag_spec.num_tags == 0):
            continue
        name = tag_spec.tag_name
        _input = text_transformed
        if (entity_tags is not None):
            _tags = self._tf_layers[f'embed.{name}.tags'](entity_tags)
            _input = tf.concat([_input, _tags], axis=(- 1))
        _logits = self._tf_layers[f'embed.{name}.logits'](_input)
        (pred_ids, confidences) = self._tf_layers[f'crf.{name}'](_logits, sequence_feature_lengths)
        predictions[f'e_{name}_ids'] = pred_ids
        predictions[f'e_{name}_scores'] = confidences
        if (name == ENTITY_ATTRIBUTE_TYPE):
            entity_tags = tf.one_hot(tf.cast(pred_ids, tf.int32), depth=tag_spec.num_tags)
    return predictions