Exemplo n.º 1
0
    def test_override_module(self):
        args = set_indexer_parser().parse_args([
            '--yaml_path',
            os.path.join(self.dir_path, 'contrib', 'fake_faiss.yml'),
            '--py_path',
            os.path.join(self.dir_path, 'contrib', 'fake_faiss.py'),
        ])

        with ServiceManager(IndexerService, args):
            pass
Exemplo n.º 2
0
    def test_override_module(self):
        args = set_indexer_parser().parse_args([
            '--yaml_path',
            os.path.join(self.dir_path, 'contrib', 'fake_faiss.yml'),
            '--py_path',
            os.path.join(self.dir_path, 'contrib', 'fake_faiss.py'),
        ])

        with ServiceManager(IndexerService, args):
            pass
        self.assertTrue(os.path.exists('foo_contrib_encoder.bin'))
        os.remove('foo_contrib_encoder.bin')
Exemplo n.º 3
0
    def test_pymode(self):
        os.unsetenv('http_proxy')
        os.unsetenv('https_proxy')
        args = set_frontend_parser().parse_args([])

        p_args = set_preprocessor_parser().parse_args([
            '--port_in',
            str(args.port_out), '--port_out', '5531', '--socket_in',
            str(SocketType.PULL_CONNECT), '--socket_out',
            str(SocketType.PUSH_BIND), '--yaml_path', 'SentSplitPreprocessor'
        ])

        e_args = set_indexer_parser().parse_args([
            '--port_in',
            str(p_args.port_out),
            '--port_out',
            str(args.port_in),
            '--socket_in',
            str(SocketType.PULL_CONNECT),
            '--socket_out',
            str(SocketType.PUSH_CONNECT),
            '--yaml_path',
            '!DictIndexer {gnes_config: {name: dummy_dict_indexer}}',
        ])

        with ServiceManager(IndexerService, e_args), \
             ServiceManager(PreprocessorService, p_args), \
             FrontendService(args), \
             grpc.insecure_channel('%s:%s' % (args.grpc_host, args.grpc_port),
                                   options=[('grpc.max_send_message_length', 70 * 1024 * 1024),
                                            ('grpc.max_receive_message_length', 70 * 1024 * 1024)]) as channel:
            stub = gnes_pb2_grpc.GnesRPCStub(channel)
            all_bytes = []
            with open(os.path.join(self.dirname, '26-doc-chinese.txt'),
                      'r',
                      encoding='utf8') as fp:
                for v in fp:
                    if v.strip():
                        all_bytes.append(v.encode())
            for r in stub.StreamCall(RequestGenerator.index(all_bytes)):
                print(r)

        bi = BaseIndexer.load('dummy_dict_indexer.bin')
        self.assertEqual(bi.size, 26)
        print(bi.query([0]))
Exemplo n.º 4
0
    def test_empty_service(self):
        args = set_indexer_parser().parse_args([
            '--yaml_path',
            '!BaseChunkIndexer {gnes_config: {name: IndexerService}}'
        ])
        c_args = _set_client_parser().parse_args(
            ['--port_in',
             str(args.port_out), '--port_out',
             str(args.port_in)])

        with ServiceManager(IndexerService, args), ZmqClient(c_args) as client:
            msg = gnes_pb2.Message()
            d = msg.request.index.docs.add()

            c = d.chunks.add()
            c.doc_id = 0
            c.embedding.CopyFrom(array2blob(self.test_numeric))
            c.offset = 0
            c.weight = 1.0

            client.send_message(msg)
            r = client.recv_message()
            self.assertEqual(r.response.index.status,
                             gnes_pb2.Response.SUCCESS)