Tensorflow: создайте input_fn для предварительной обработки текста с помощью API tf.data.Dataset

Question

Я работаю через обучение модели Tensorflow и использую tf.data.Dataset API. Я работаю через этот урок:

https://nbviewer.jupyter.org/github/amygdala/code-snippets/blob/master/ml/census_train_and_eval/using_tf.estimator.train_and_evaluate.ipynb

Короче говоря, у меня есть панды DataFrame, dfс колонкой под названием text и столбец под названием label, Я хотел бы создать input_fn который возвращает кортеж (feature, indexes), где features - это словарь Tensors, а index - это единственный тензор индексов меток.

Все, что я хотел бы сделать, это разбить мою строку на пробел, преобразовать слова в целые числа (каждое слово получает уникальное целое число) и заполнить до длины 100. Это версия учебника:

def input_fn(filenames,
                      num_epochs=None,
                      shuffle=True,
                      skip_header_lines=0,
                      batch_size=200):
  """Generates features and labels for training or evaluation.
  """

  dataset = tf.data.TextLineDataset(filenames).skip(skip_header_lines).map(parse_csv)

  if shuffle:
    dataset = dataset.shuffle(buffer_size=batch_size * 10)
  dataset = dataset.repeat(num_epochs)
  dataset = dataset.batch(batch_size)
  iterator = dataset.make_one_shot_iterator()
  features = iterator.get_next()
  return features, parse_label_column(features.pop(LABEL_COLUMN))

Я думал начать с чего-то вроде:

training_dataset = (
    tf.data.Dataset.from_tensor_slices(
        (
            tf.cast(df["text"].values, tf.float32),
            tf.cast(df["label"].values, tf.int32)
        )
    )
)

Но я немного не уверен, как поступить. Я думал:

def pad_or_trunc(t):
    dim = tf.size(t)
    return tf.cond(tf.equal(dim, k), lambda: t, lambda: tf.cond(tf.greater(dim, k), lambda: tf.slice(t, [0], [k]), lambda: tf.concat([t, tf.zeros(k-dim, dtype=tf.int32)], 0)))

training_dataset = training_dataset.map(lambda string: tf.string_split([string]).values)
training_dataset = training_dataset.map(tf.feature_column.categorical_column_with_identity('text', vocab_size=1000))

training_dataset = training_dataset.map(pad_or_trunc)

Я чувствую, что я изобретаю колесо здесь, но документы по использованию tf.data.Dataset API с текстом довольно редки... есть идеи по поводу лучшего подхода?