Функция прогнозирования предикторов для классификации текста с использованием сохраненной модели BERT

Question

Я создал модель BERT для классификации сгенерированной пользователем текстовой строки как FAQ или не FAQ. Я сохранил свою модель, используя функцию export_savedmodel (). Я хочу написать функцию, чтобы предсказать вывод для нового набора строк, который принимает в качестве входных данных список строк.

Я пытался использовать метод преломитель.from_saved_model (), но этот метод требует передачи пар ключ-значение для ввода идентификатора, идентификатора сегмента, идентификатора метки и маски ввода. Я новичок и не совсем понимаю, что здесь проходить.

Экспорт или сохранение модели

def serving_input_fn():
    label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
    input_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='input_ids')
    input_mask = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='input_mask')
    segment_ids = tf.placeholder(tf.int32, [None, MAX_SEQ_LENGTH], name='segment_ids')
    input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
        'label_ids': label_ids,
        'input_ids': input_ids,
        'input_mask': input_mask,
        'segment_ids': segment_ids,
    })()
    return input_fn

export_dir = "..."
estimator._export_to_tpu = False
estimator.export_savedmodel(export_dir, serving_input_fn)

#Predicting
with tf.Session() as sess:   
    predict_fn = predictor.from_saved_model(...')

#Data description
My data is a simple table having a column for input string and another for output label.

# Error.
ValueError: Got unexpected keys in input_dict: {'pred'}
expected: {'label_ids', 'input_mask', 'segment_ids', 'input_ids'}

#Thank you for any help!