ngoanpv / albert_vi Goto Github PK

import sentencepiece as spm
# import tensorflow_hub as hub
import tensorflow as tf

sp = spm.SentencePieceProcessor(model_file='assets/albertvi_30k-clean.model')

input_ids = sp.encode("Tui rất là đẹp trai.")


tags = set()
# if is_training:
#   tags.add("train")
# albert_module = hub.Module("./base", tags=tags,
#                            trainable=False)
with tf.Session() as sess:
    saver = tf.compat.v1.train.import_meta_graph('base/model.ckpt-1015000.meta')
    saver.restore(sess, 'base/model.ckpt-1015000.data-00000-of-00001')

    # graph = tf.get_default_graph()

    albert_inputs = dict(
        input_ids=input_ids,
        # input_mask=input_mask,
        # segment_ids=segment_ids
    )
    # albert_outputs = model(
    #     inputs=albert_inputs,
    #     signature="tokens",
    #     as_dict=True)

# If you want to use the token-level output, use
# albert_outputs["sequence_output"] instead.
#     output_layer = albert_outputs["pooled_output"]