def do_write_callback(self, bsz): d, n = 32, 1000 x = np.random.uniform(size=(n, d)).astype('float32') index = faiss.IndexFlatL2(d) index.add(x) f = io.BytesIO() # test with small block size writer = faiss.PyCallbackIOWriter(f.write, 1234) if bsz > 0: writer = faiss.BufferedIOWriter(writer, bsz) faiss.write_index(index, writer) del writer # make sure all writes committed if sys.version_info[0] < 3: buf = f.getvalue() else: buf = f.getbuffer() index2 = faiss.deserialize_index(np.frombuffer(buf, dtype='uint8')) self.assertEqual(index.d, index2.d) self.assertTrue( np.all( faiss.vector_to_array(index.xb) == faiss.vector_to_array( index2.xb))) # This is not a callable function: shoudl raise an exception writer = faiss.PyCallbackIOWriter("blabla") self.assertRaises(Exception, faiss.write_index, index, writer)
def test_transfer_pipe(self): """ transfer an index through a Unix pipe """ d, n = 32, 1000 x = np.random.uniform(size=(n, d)).astype('float32') index = faiss.IndexFlatL2(d) index.add(x) Dref, Iref = index.search(x, 10) rf, wf = os.pipe() # start thread that will decompress the index def index_from_pipe(): reader = faiss.PyCallbackIOReader(lambda size: os.read(rf, size)) return faiss.read_index(reader) with ThreadPool(1) as pool: fut = pool.apply_async(index_from_pipe, ()) # write to pipe writer = faiss.PyCallbackIOWriter(lambda b: os.write(wf, b)) faiss.write_index(index, writer) index2 = fut.get() # closing is not really useful but it does not hurt os.close(wf) os.close(rf) Dnew, Inew = index2.search(x, 10) np.testing.assert_array_equal(Iref, Inew) np.testing.assert_array_equal(Dref, Dnew)
def output_faiss_index(self): print('faiss index ntotal [%d]: %d' % (self.rank, self.faiss_index.ntotal)) item_index_stream = _mindalpha.OutputStream( self.item_index_output_path) item_index_writer = faiss.PyCallbackIOWriter(item_index_stream.write) faiss.write_index(self.faiss_index, item_index_writer) self.output_index_meta() self.item_ids_stream = None