import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from google.protobuf import text_format
import os

export_dir = 'inference/pb2saved'
graph_pb = '/Users/edz/PycharmProjects/serviceRec/train/saved_model_test/1640591747/saved_model.pb'

if os.path.exists(export_dir):
    os.rmdir(export_dir)

builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

with tf.gfile.GFile(graph_pb, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

sigs = {}

with tf.Session(graph=tf.Graph()) as sess:
    # name="" is important to ensure we don't get spurious prefixing
    tf.import_graph_def(graph_def, name="")
    g = tf.get_default_graph()

    print(sess.graph.get_name_scope())
    print(sess.graph.get_all_collection_keys())
    print(sess.graph.get_operations())

    # input_ids = sess.graph.get_tensor_by_name(
    # "input_ids:0")
    # input_mask = sess.graph.get_tensor_by_name(
    # "input_mask:0")
    # segment_ids = sess.graph.get_tensor_by_name(
    # "segment_ids:0")
    # probabilities = g.get_tensor_by_name("loss/pred_prob:0")

    # sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
    #     tf.saved_model.signature_def_utils.predict_signature_def(
    #         {
    #             "input_ids": input_ids,
    #             "input_mask": input_mask,
    #             "segment_ids": segment_ids
    #         }, {
    #             "probabilities": probabilities
    #             })

    # builder.add_meta_graph_and_variables(sess,
    #                                      [tag_constants.SERVING],
    #                                      signature_def_map=sigs)

# builder.save()