import tensorflow as tf

            "with_src_lang_tag": self._with_src_lang_tag,
            "trg_lang_tag_position": self._trg_lang_tag_position,
        }

    def inputs_signature(self, mode):
        """ Returns the input dtypes and signatures. """
        dtypes = {"feature": tf.int64, "src_lang": tf.int64, "trg_lang": tf.int64}
        signatures = {"feature": tf.TensorShape([None, None]),
                      "src_lang": tf.TensorShape([None, ]),
                      "trg_lang": tf.TensorShape([None, ])}
        if mode == compat.ModeKeys.INFER:
            return dtypes, signatures
        dtypes["label"] = tf.int64
        signatures["label"] = tf.TensorShape([None, None])
        return dtypes, signatures

