记java重构python版bert-serving-client

背景

项目需要把bert-serving-client由python用java实现，因为java比python快一些，于是就开始了尝试

先上bert-as-service的github地址：https://github.com/hanxiao/bert-as-service

其中client的init.py文件地址：https://github.com/hanxiao/bert-as-service/blob/master/client/bert_serving/client/__init__.py

主要实现其中encode、fetch、fetchAll和encodeAsync

导包

bertClient主要用到zeroMq和json，前者用来提供和服务端的连接，后者格式化传输数据。两者pom依赖如下

        
            org.zeromq
            jeromq
            0.5.1
        

        
        
            org.zeromq
            jeromq
            0.5.2-SNAPSHOT
        

        
            com.google.code.gson
            gson
            2.8.2
        

        
            org.json
            json
            20180813
        

构造方法

python中有默认参数，java里没有，于是我采取属性的默认值+方法重载来实现默认参数。最后java版的构造函数如下：

    private void init() throws Exception {
        mContext = new ZContext();
        String url = "tcp://" + mIp + ":";
        mIdentity = UUID.randomUUID().toString();

        mSendSocket = mContext.createSocket(SocketType.PUSH);
        mSendSocket.setLinger(0);
        mSendSocket.connect(url + mPort);

        mRecvSocket = mContext.createSocket(SocketType.SUB);
        mRecvSocket.setLinger(0);
        mRecvSocket.subscribe(mIdentity.getBytes(CHARSET_NAME));
        mRecvSocket.connect(url + mPortOut);
    }

对应python版的构造函数：

    def __init__(self, ip='localhost', port=5555, port_out=5556,
                 output_fmt='ndarray', show_server_config=False,
                 identity=None, check_version=True, check_length=True,
                 check_token_info=True, ignore_all_checks=False,
                 timeout=-1):

        self.context = zmq.Context()
        self.sender = self.context.socket(zmq.PUSH)
        self.sender.setsockopt(zmq.LINGER, 0)
        self.identity = identity or str(uuid.uuid4()).encode('ascii')
        self.sender.connect('tcp://%s:%d' % (ip, port))

        self.receiver = self.context.socket(zmq.SUB)
        self.receiver.setsockopt(zmq.LINGER, 0)
        self.receiver.setsockopt(zmq.SUBSCRIBE, self.identity)
        self.receiver.connect('tcp://%s:%d' % (ip, port_out))

        ....

        ....

收发数据

收发数据对应python版里的_send()和_recv()函数，两者代码如下

    def _send(self, msg, msg_len=0):
        self.request_id += 1
        self.sender.send_multipart([self.identity, msg, b'%d' % self.request_id, b'%d' % msg_len])
        self.pending_request.add(self.request_id)
        return self.request_id

    def _recv(self, wait_for_req_id=None):
        try:
            while True:
                # a request has been returned and found in pending_response
                if wait_for_req_id in self.pending_response:
                    response = self.pending_response.pop(wait_for_req_id)
                    return _Response(wait_for_req_id, response)

                # receive a response
                response = self.receiver.recv_multipart()
                request_id = int(response[-1])

                # if not wait for particular response then simply return
                if not wait_for_req_id or (wait_for_req_id == request_id):
                    self.pending_request.remove(request_id)
                    return _Response(request_id, response)
                elif wait_for_req_id != request_id:
                    self.pending_response[request_id] = response
                    # wait for the next response
        except Exception as e:
            raise e
        finally:
            if wait_for_req_id in self.pending_request:
                self.pending_request.remove(wait_for_req_id)

_send()函数里主要调用了发送套接字的send_multipart()函数，把identity、msg、request_id和msg_len作为列表发送过去，java里没有直接对应send_multipart()的方法，可以用sendMore()和send()代替

同样，_recv()函数里主要调用了接收套接字的recv_multipart()函数，java中也没有直接对应的方法，可以用recvMore()代替，最后可以写出java版代码如下

    public long send(String message) {
        return send(message, 0);
    }

    public long send(String message, int messageLen) {
        return send(new String[]{message}, messageLen);
    }

    public long send(String[] message, int messageLen) {
        mRequestId++;
        Gson gson = new Gson();
        sendMultiPart(new String[]{mIdentity, gson.toJson(message), mRequestId + "", messageLen + ""});
        mPendingRequest.add(mRequestId);
        return mRequestId;
    }

    private void sendMultiPart(String[] msgParts) {
        try {
            int i;
            for (i = 0; i < msgParts.length - 1; i++) {
                mSendSocket.sendMore(msgParts[i].getBytes(CHARSET_NAME));
            }
            mSendSocket.send(msgParts[i].getBytes(CHARSET_NAME), 0);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public Map recv() {
        return recv(null);
    }

    public Map recv(Long waitForReqId) {
        try {
            while (true) {
                if (waitForReqId != null && mPendingResponse.containsKey(waitForReqId)) {
                    List response = mPendingResponse.get(waitForReqId);
                    HashMap resultMap = new HashMap();
                    resultMap.put(KEY_ID, waitForReqId);
                    resultMap.put(KEY_CONTENT, response);
                    return resultMap;
                }

                List response = recvMutipart(0);
                if (response == null || response.size() == 0) {
                    return null;
                }

                long requestId = Utils.byte2Long(response.get(response.size() - 1));


                if (waitForReqId == null || waitForReqId == requestId) {
                    mPendingRequest.remove(requestId);
                    HashMap resultMap = new HashMap();
                    resultMap.put(KEY_ID, requestId);
                    if (response != null) {
                        resultMap.put(KEY_CONTENT, response);
                    }
                    return resultMap;
                } else if (waitForReqId != requestId) {
                    mPendingResponse.put(requestId, response);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (waitForReqId != null && mPendingRequest.contains(waitForReqId)) {
                mPendingRequest.remove(waitForReqId);
            }
        }
        return null;
    }

    private List recvMutipart(int flag) {
        ArrayList result = new ArrayList();
        byte[] item = mRecvSocket.recv(flag);
        if (item != null) {
            result.add(item);
        }
        while (mRecvSocket.hasReceiveMore()) {
            item = mRecvSocket.recv(flag);
            if (item != null) {
                result.add(item);
            }
        }
        return result;
    }

注意send()方法中，发送消息时，一定要用gson把消息字符串转换成json格式，否则服务端会报错，客户端收不到数据

自定义的sendMultiPart()方法中，把字符串编码成字节数组时用的编码格式是utf-8，用了自定义常量显示

在接收回复时，根据python版的代码可知，每一个回复的最后一部分是发送消息时对应的请求id，可以采取措施把byte[]数组转换成Long，具体代码如下

    public static long byte2Long(byte[] bytes) {
        if (bytes == null) {
            return -1L;
        }
        StringBuilder builder = new StringBuilder();
        for (int i = 0; i < bytes.length; i++) {
            builder.append(bytes[i] - 48);
        }
        return Long.parseLong(builder.toString());
    }

比如收到的字节数组是[49, 48]，显然对应的请求id是10，那么由上面的byte2Long方法就可以进行转换

encode

收发数据完成之后，就可以着力实现encode、fetch、fetchAll和encodeAsync四个方法了。

encode编码字符串，在调试python版的客户端后发现，encode编码的字符串必须首先转换成字符串数组，比如"szc"要转换成["s", "z", "c"]。根据这一点，以及python版代码，可以写出java版encode()方法和重载方法

    public List