Как сформировать входные данные для использования вложенных слов BPEmb в модели Keras LSTM для генерации текста?

Question

Я пытаюсь обучить модель Keras / Tensorflow LSTM в Колаборатории Google для генерации текста из подслов с использованием вложений BPEmb .

Насколько я понимаю, вложения - это модель генезиса, которая производит 300-мерные векторы для кусков подслов.

Моя модель предназначена для использования последовательностей k вложений для прогнозирования k+1 -го внедрения, которое я затем декодирую, используя BPEmb.similar_by_vector(), а затем BPEmb.decode().

У меня есть ~ 2,5M-встраиваемый набор данных, поэтому для длины последовательности 100 мне нужно обработать 2.5 * 10^8 * 300 значения. Чтобы избежать переполнения памяти, я использую функцию генератора следующим образом:

# this reads my data into a huge text file
text = read_data(DATA_FILES) 
bpe_uk = BPEmb(lang='uk', dim=EMB_DIM)
vectors = bpe_uk.embed(text)

def training_generator2(vectors, k=100, batch_size=1024):
    """Dynamically feed training sequences for a recurrent neural network model:
    Output shape:
    x = (batch_size, k, dim)
    y = (300)
    """

    i = 0  # offset in the data set
    while i+k < len(vectors):

        # init empty batch
        x = []
        for sample_id in range(0, batch_size):

          x.append(vectors[i:i+k])
          y = vectors[i+k]
          i+=1
        yield (x, y)

g = training_generator2(vectors, 100, 128)
v = next(g)

np.shape(v[0])
>>> (128, 100, 300) # the X of the training pair is 128 batches of 100 sequences of 300-dim vectors

np.shape(v[1])
>>> (300, ) # the Y of the training pair is just one 300-dim vector

Мое определение модели выглядит следующим образом:

EMBEDDING_DIM = EMB_DIM
SHAPE = (SEQ_LEN, EMB_DIM)

def lstm_model(seq_len=SEQ_LEN, batch_size=None, stateful=True):
  """Language model: predict the next word given the current word."""
  source = tf.keras.Input(
      name='seed', shape=(seq_len, EMBEDDING_DIM), batch_size=batch_size, dtype=tf.float32)

  lstm_1 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(source)
  lstm_2 = tf.keras.layers.LSTM(EMBEDDING_DIM, stateful=stateful, return_sequences=True)(lstm_1)
  predicted_char = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(300, activation='softmax'))(lstm_2)
  model = tf.keras.Model(inputs=[source], outputs=[predicted_char])
  model.compile(
      optimizer=tf.train.RMSPropOptimizer(learning_rate=0.01),
      loss='sparse_categorical_crossentropy',
      metrics=['sparse_categorical_accuracy'])
  return model

model = lstm_model(seq_len=SEQ_LEN, batch_size=128, stateful=False)
model.summary()

Layer (type)                 Output Shape              Param #   
=================================================================
seed (InputLayer)            (128, 100, 300)           0         
_________________________________________________________________
lstm_3 (LSTM)                (128, 100, 300)           721200    
_________________________________________________________________
lstm_4 (LSTM)                (128, 100, 300)           721200    
_________________________________________________________________
time_distributed_1 (TimeDist (128, 100, 300)           90300     
=================================================================
Total params: 1,532,700
Trainable params: 1,532,700
Non-trainable params: 0
_________________________________________________________________

Когда я пытаюсь тренировать эту модель так:

tf.keras.backend.clear_session()

training_model = lstm_model(seq_len=SEQ_LEN, batch_size=128, stateful=False)

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    training_model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(TPU_WORKER)))

tpu_model.fit_generator(
    training_generator2(vectors, k=100, batch_size=128),
    steps_per_epoch=100,
    epochs=10,
)

Я получаю следующее сообщение об ошибке:

ValueError: Error when checking input: expected seed to have 3 dimensions, but got array with shape (100, 300)

Похоже, данные обучения должны быть преобразованы во что-то вроде (100, 300, 1), причем окончательные значения в основном являются массивами одного значения. Это правильно? И если да, то как это сделать?

Обратите внимание, что я новичок в Keras и в целом глубоко изучаю, поэтому может возникнуть несколько проблем с этой настройкой, включая определение модели. Любые отзывы или предложения будут приветствоваться!