bert serving

BertModel.java

import org.apache.commons.io.FileUtils;
import org.tensorflow.Graph;
import org.tensorflow.SavedModelBundle;
import org.tensorflow.Session;
import org.tensorflow.Tensor;

import java.io.File;
import java.io.IOException;
import java.util.List;

// check out exported model: saved_model_cli show --all --dir $exported_dir

public class BertModel {

    public static void main(String[] args) throws IOException {
//        try (Graph graph = new Graph()) {
//            graph.importGraphDef(FileUtils.readFileToByteArray(new File(
//                    "E:\\dlprojects\\data\\bert\\chinese_L-12_H-768_A-12\\pb\\1561992264\\saved_model.pb")));
//
//        }
        SavedModelBundle bundle=SavedModelBundle.load("E:\\dlprojects\\data\\bert\\chinese_L-12_H-768_A-12\\pb\\1561992264","serve");
        Session s = bundle.session();

        int seqLen = 128;
        int batchSize = 64;
        int[] unique_ids = new int[batchSize];
        int[][] input_ids = new int[batchSize][seqLen];
        int[][] input_mask = new int[batchSize][seqLen];
        int[][] input_type_ids = new int[batchSize][seqLen];

        Tensor uniqueIdsTensor = Tensor.create(unique_ids);
        Tensor inputIdsTensor = Tensor.create(input_ids);
        Tensor inputMaskTensor = Tensor.create(input_mask);
        Tensor inputTypeIdsTensor = Tensor.create(input_type_ids);

        Long a = System.currentTimeMillis();
        System.out.println(s.runner()
                .feed("unique_ids_1", uniqueIdsTensor)
                .feed("input_ids_1", inputIdsTensor)
                .feed("input_mask_1", inputMaskTensor)
                .feed("input_type_ids_1", inputTypeIdsTensor)
                .fetch("bert/encoder/Reshape_13:0")
                .run());
        Long b = System.currentTimeMillis();
        System.out.println(b-a);

    }
}