import tensorflow as tf


class TranslateDistillProblem(TranslateProblem):
  """Base class for translation problems."""

  def is_generate_per_split(self):
    return True

  def example_reading_spec(self):
    data_fields = {"dist_targets": tf.VarLenFeature(tf.int64)}

    if self.has_inputs:
      data_fields["inputs"] = tf.VarLenFeature(tf.int64)

    # hack: ignoring true targets and putting dist_targets in targets
    data_items_to_decoders = {
        "inputs": tf.contrib.slim.tfexample_decoder.Tensor("inputs"),
