-
Notifications
You must be signed in to change notification settings - Fork 306
/
faiss_index_pai.py
112 lines (96 loc) · 3.95 KB
/
faiss_index_pai.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
from __future__ import print_function
import logging
import os
import faiss
import numpy as np
import tensorflow as tf
logging.basicConfig(
level=logging.INFO, format='[%(asctime)s][%(levelname)s] %(message)s')
tf.app.flags.DEFINE_string('tables', '', 'tables passed by pai command')
tf.app.flags.DEFINE_integer('batch_size', 1024, 'batch size')
tf.app.flags.DEFINE_integer('embedding_dim', 32, 'embedding dimension')
tf.app.flags.DEFINE_string('index_output_dir', '', 'index output directory')
tf.app.flags.DEFINE_string('index_type', 'IVFFlat', 'index type')
tf.app.flags.DEFINE_integer('ivf_nlist', 1000, 'nlist')
tf.app.flags.DEFINE_integer('hnsw_M', 32, 'hnsw M')
tf.app.flags.DEFINE_integer('hnsw_efConstruction', 200, 'hnsw efConstruction')
tf.app.flags.DEFINE_integer('debug', 0, 'debug index')
FLAGS = tf.app.flags.FLAGS
def main(argv):
reader = tf.python_io.TableReader(
FLAGS.tables, slice_id=0, slice_count=1, capacity=FLAGS.batch_size * 2)
i = 0
id_map_f = tf.gfile.GFile(
os.path.join(FLAGS.index_output_dir, 'id_mapping'), 'w')
embeddings = []
while True:
try:
records = reader.read(FLAGS.batch_size)
for j, record in enumerate(records):
if isinstance(record[0], bytes):
eid = record[0].decode('utf-8')
id_map_f.write('%s\n' % eid)
embeddings.extend(
[list(map(float, record[1].split(b','))) for record in records])
i += 1
if i % 100 == 0:
logging.info('read %d embeddings.' % (i * FLAGS.batch_size))
except tf.python_io.OutOfRangeException:
break
reader.close()
id_map_f.close()
logging.info('Building faiss index..')
if FLAGS.index_type == 'IVFFlat':
quantizer = faiss.IndexFlatIP(FLAGS.embedding_dim)
index = faiss.IndexIVFFlat(quantizer, FLAGS.embedding_dim, FLAGS.ivf_nlist,
faiss.METRIC_INNER_PRODUCT)
elif FLAGS.index_type == 'HNSWFlat':
index = faiss.IndexHNSWFlat(FLAGS.embedding_dim, FLAGS.hnsw_M,
faiss.METRIC_INNER_PRODUCT)
index.hnsw.efConstruction = FLAGS.hnsw_efConstruction
else:
raise NotImplementedError
embeddings = np.array(embeddings)
if FLAGS.index_type == 'IVFFlat':
logging.info('train embeddings...')
index.train(embeddings)
logging.info('build embeddings...')
index.add(embeddings)
faiss.write_index(index, 'faiss_index')
with tf.gfile.GFile(
os.path.join(FLAGS.index_output_dir, 'faiss_index'), 'wb') as f_out:
with open('faiss_index', 'rb') as f_in:
f_out.write(f_in.read())
if FLAGS.debug != 0:
# IVFFlat
for ivf_nlist in [100, 500, 1000, 2000]:
quantizer = faiss.IndexFlatIP(FLAGS.embedding_dim)
index = faiss.IndexIVFFlat(quantizer, FLAGS.embedding_dim, ivf_nlist,
faiss.METRIC_INNER_PRODUCT)
index.train(embeddings)
index.add(embeddings)
index_name = 'faiss_index_ivfflat_nlist%d' % ivf_nlist
faiss.write_index(index, index_name)
with tf.gfile.GFile(
os.path.join(FLAGS.index_output_dir, index_name), 'wb') as f_out:
with open(index_name, 'rb') as f_in:
f_out.write(f_in.read())
# HNSWFlat
for hnsw_M in [16, 32, 64, 128]:
for hnsw_efConstruction in [64, 128, 256, 512, 1024, 2048, 4096, 8196]:
if hnsw_efConstruction < hnsw_M * 2:
continue
index = faiss.IndexHNSWFlat(FLAGS.embedding_dim, hnsw_M,
faiss.METRIC_INNER_PRODUCT)
index.hnsw.efConstruction = hnsw_efConstruction
index.add(embeddings)
index_name = 'faiss_index_hnsw_M%d_ef%d' % (hnsw_M, hnsw_efConstruction)
faiss.write_index(index, index_name)
with tf.gfile.GFile(
os.path.join(FLAGS.index_output_dir, index_name), 'wb') as f_out:
with open(index_name, 'rb') as f_in:
f_out.write(f_in.read())
if __name__ == '__main__':
tf.app.run()