-
Notifications
You must be signed in to change notification settings - Fork 15
/
run_norm_rq.py
30 lines (26 loc) · 935 Bytes
/
run_norm_rq.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
from vecs_io import *
from pq_residual import *
from sorter import *
from run_pq import execute
from run_pq import parse_args
if __name__ == '__main__':
dataset = 'netflix'
topk = 20
codebook = 4
Ks = 256
metric = 'product'
# override default parameters with command line parameters
import sys
if len(sys.argv) > 3:
dataset, topk, codebook, Ks, metric = parse_args()
else:
import warnings
warnings.warn("Using Default Parameters ")
print("# Parameters: dataset = {}, topK = {}, codebook = {}, Ks = {}, metric = {}"
.format(dataset, topk, codebook, Ks, metric))
X, T, Q, G = loader(dataset, topk, metric, folder='data/')
# pq, rq, or component of norm-pq
pqs = [PQ(M=1, Ks=Ks) for _ in range(codebook-1)]
quantizer = ResidualPQ(pqs=pqs)
quantizer = NormPQ(n_percentile=Ks, quantize=quantizer)
execute(quantizer, X, T, Q, G, metric)