<s> """<STR_LIT>""" <EOL> DESCRIPTION = """<STR_LIT:U+0020>""" <EOL> import logging <EOL> import numpy as np <EOL> from hdidx . encoder import Encoder <EOL> from hdidx . util import kmeans , pq_kmeans_assign <EOL> class PQEncoder ( Encoder ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> super ( PQEncoder , self ) . __init__ ( ) <EOL> def __del__ ( self ) : <EOL> pass <EOL> def build ( self , pardic = None ) : <EOL> vals = pardic [ '<STR_LIT>' ] <EOL> nsubq = pardic [ '<STR_LIT>' ] <EOL> nsubqbits = pardic . get ( '<STR_LIT>' , <NUM_LIT:8> ) <EOL> blksize = pardic . get ( '<STR_LIT>' , <NUM_LIT> ) <EOL> dim = vals . shape [ <NUM_LIT:1> ] <EOL> dsub = dim / nsubq <EOL> ksub = <NUM_LIT:2> ** nsubqbits <EOL> """<STR_LIT>""" <EOL> ecdat = dict ( ) <EOL> ecdat [ '<STR_LIT>' ] = nsubq <EOL> ecdat [ '<STR_LIT>' ] = ksub <EOL> ecdat [ '<STR_LIT>' ] = dsub <EOL> ecdat [ '<STR_LIT>' ] = blksize <EOL> ecdat [ '<STR_LIT>' ] = [ None for q in range ( nsubq ) ] <EOL> logging . info ( "<STR_LIT>" ) <EOL> for q in range ( nsubq ) : <EOL> logging . info ( "<STR_LIT>" % ( q , nsubq ) ) <EOL> vs = np . require ( vals [ : , q * dsub : ( q + <NUM_LIT:1> ) * dsub ] , <EOL> requirements = '<STR_LIT:C>' , dtype = np . float32 ) <EOL> ecdat [ '<STR_LIT>' ] [ q ] = kmeans ( vs , ksub , niter = <NUM_LIT:100> ) <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . ecdat = ecdat <EOL> def encode ( self , vals ) : <EOL> dsub = self . ecdat [ '<STR_LIT>' ] <EOL> nsubq = self . ecdat [ '<STR_LIT>' ] <EOL> centroids = self . ecdat [ '<STR_LIT>' ] <EOL> num_vals = vals . shape [ <NUM_LIT:0> ] <EOL> codes = np . zeros ( ( num_vals , nsubq ) , np . uint8 ) <EOL> for q in range ( nsubq ) : <EOL> vsub = vals [ : , q * dsub : ( q + <NUM_LIT:1> ) * dsub ] <EOL> codes [ : , q ] = pq_kmeans_assign ( centroids [ q ] , vsub ) <EOL> return codes <EOL> class IVFPQEncoder ( PQEncoder ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> super ( IVFPQEncoder , self ) . __init__ ( ) <EOL> def __del__ ( self ) : <EOL> pass <EOL> def build ( self , pardic = None ) : <EOL> pardic [ '<STR_LIT>' ] = pardic [ '<STR_LIT>' ] . copy ( ) <EOL> vals = pardic [ '<STR_LIT>' ] <EOL> coarsek = pardic . get ( '<STR_LIT>' , <NUM_LIT> ) <EOL> logging . info ( '<STR_LIT>' ) <EOL> coa_centroids = kmeans ( vals . astype ( np . float32 ) , coarsek , niter = <NUM_LIT:100> ) <EOL> cids = pq_kmeans_assign ( coa_centroids , vals ) <EOL> logging . info ( '<STR_LIT>' ) <EOL> pardic [ '<STR_LIT>' ] -= coa_centroids [ cids , : ] <EOL> super ( IVFPQEncoder , self ) . build ( pardic ) <EOL> self . ecdat [ '<STR_LIT>' ] = coa_centroids <EOL> self . ecdat [ '<STR_LIT>' ] = coarsek <EOL> def encode ( self , vals ) : <EOL> vals = vals . copy ( ) <EOL> coa_centroids = self . ecdat [ '<STR_LIT>' ] <EOL> cids = pq_kmeans_assign ( coa_centroids , vals ) <EOL> vals -= coa_centroids [ cids , : ] <EOL> codes = super ( IVFPQEncoder , self ) . encode ( vals ) <EOL> return cids , codes </s>
