Как предсказать последовательность букв с помощью tensorflow?

Задача: предсказать следующую букву в твите, имея заданную начальную последовательность

X = pad_sequences(data, maxlen=149, padding='post', truncating='post')
Y = np.array([np.append(i[1:], 0) for i in X])

Я использую модель тензорфлоу и отправляю туда две матрицы Х и у, последовательность букв в твитах и эту же последовательность смещенную на букву вперед. Всего возможных символов - 64

model = tf.keras.Sequential()
model.add(tf.keras.layers.Embedding(input_dim=64, output_dim=8, 
                                    mask_zero=True, input_length=149))
model.add(tf.keras.layers.LSTM(128, return_sequences=True))
model.add(tf.keras.layers.Dense(64, activation='softmax'))
model.compile(optimizer='Adam', loss='sparse_categorical_crossentropy')
history = model.fit(X,Y, epochs=5, batch_size=16, validation_split=0.2)

Но дальше я не совсем понимаю, что мне использовать. Я получаю на выход столько же векторов длиной 64, сколько букв было в заданной последовательности (seq), каждый вектор - это полученные вероятности на один шаг последовательности? Я пробовал три варианта, все они не дают идеальный результат. Не могу найти ошибку.

pred = model.predict(seq)

a1=pred[-1].argsort().ravel()[::-1] # последний вектор, сортировка по вероятности
pred_chars1 = ['' if i==0 else tokenizer.index_word[i] for i in a1]
tw1 = pred_chars1[0]

a2 = pred.sum(axis=0).argsort(axis=1).ravel()[::-1] # сумма весов по всем векторам, сортировка по вероятности
pred_chars2 = ['' if i==0 else tokenizer.index_word[i] for i in a2]
tw2 = pred_chars2[0]

a3 = pred.argmax(axis=0).ravel()[::-1] # наибольшая вероятность для каждого вектора
pred_chars3 = ['' if i==0 else tokenizer.index_word[i] for i in a3]
tw3 = pred_chars3[0]