<s> """<STR_LIT>""" <EOL> import logging <EOL> import itertools <EOL> import os <EOL> import heapq <EOL> import numpy <EOL> import scipy . sparse <EOL> from gensim import interfaces , utils , matutils <EOL> from six . moves import map as imap , xrange , zip as izip <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> PARALLEL_SHARDS = False <EOL> try : <EOL> import multiprocessing <EOL> except ImportError : <EOL> pass <EOL> class Shard ( utils . SaveLoad ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , fname , index ) : <EOL> self . dirname , self . fname = os . path . split ( fname ) <EOL> self . length = len ( index ) <EOL> self . cls = index . __class__ <EOL> logger . info ( "<STR_LIT>" % self . fullname ( ) ) <EOL> index . save ( self . fullname ( ) ) <EOL> self . index = self . get_index ( ) <EOL> def fullname ( self ) : <EOL> return os . path . join ( self . dirname , self . fname ) <EOL> def __len__ ( self ) : <EOL> return self . length <EOL> def __getstate__ ( self ) : <EOL> result = self . __dict__ . copy ( ) <EOL> if '<STR_LIT:index>' in result : <EOL> del result [ '<STR_LIT:index>' ] <EOL> return result <EOL> def __str__ ( self ) : <EOL> return ( "<STR_LIT>" % ( self . cls . __name__ , len ( self ) , self . fullname ( ) ) ) <EOL> def get_index ( self ) : <EOL> if not hasattr ( self , '<STR_LIT:index>' ) : <EOL> logger . debug ( "<STR_LIT>" % self . fullname ( ) ) <EOL> self . index = self . cls . load ( self . fullname ( ) , mmap = '<STR_LIT:r>' ) <EOL> return self . index <EOL> def get_document_id ( self , pos ) : <EOL> """<STR_LIT>""" <EOL> assert <NUM_LIT:0> <= pos < len ( self ) , "<STR_LIT>" <EOL> return self . get_index ( ) . index [ pos ] <EOL> def __getitem__ ( self , query ) : <EOL> index = self . get_index ( ) <EOL> try : <EOL> index . num_best = self . num_best <EOL> index . normalize = self . normalize <EOL> except : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> return index [ query ] <EOL> def query_shard ( args ) : <EOL> query , shard = args <EOL> logger . debug ( "<STR_LIT>" % ( shard , shard . num_best , os . getpid ( ) ) ) <EOL> result = shard [ query ] <EOL> logger . debug ( "<STR_LIT>" % ( shard , os . getpid ( ) ) ) <EOL> return result <EOL> class Similarity ( interfaces . SimilarityABC ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , output_prefix , corpus , num_features , num_best = None , chunksize = <NUM_LIT> , shardsize = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> if output_prefix is None : <EOL> self . output_prefix = utils . randfname ( prefix = '<STR_LIT>' ) <EOL> else : <EOL> self . output_prefix = output_prefix <EOL> logger . info ( "<STR_LIT>" % self . output_prefix ) <EOL> self . num_features = num_features <EOL> self . num_best = num_best <EOL> self . normalize = True <EOL> self . chunksize = int ( chunksize ) <EOL> self . shardsize = shardsize <EOL> self . shards = [ ] <EOL> self . fresh_docs , self . fresh_nnz = [ ] , <NUM_LIT:0> <EOL> if corpus is not None : <EOL> self . add_documents ( corpus ) <EOL> def __len__ ( self ) : <EOL> return len ( self . fresh_docs ) + sum ( [ len ( shard ) for shard in self . shards ] ) <EOL> def __str__ ( self ) : <EOL> return ( "<STR_LIT>" % <EOL> ( len ( self ) , len ( self . shards ) , self . output_prefix ) ) <EOL> def add_documents ( self , corpus ) : <EOL> """<STR_LIT>""" <EOL> min_ratio = <NUM_LIT:1.0> <EOL> if self . shards and len ( self . shards [ - <NUM_LIT:1> ] ) < min_ratio * self . shardsize : <EOL> self . reopen_shard ( ) <EOL> for doc in corpus : <EOL> if isinstance ( doc , numpy . ndarray ) : <EOL> doclen = len ( doc ) <EOL> elif scipy . sparse . issparse ( doc ) : <EOL> doclen = doc . nnz <EOL> else : <EOL> doclen = len ( doc ) <EOL> if doclen < <NUM_LIT> * self . num_features : <EOL> doc = matutils . unitvec ( matutils . corpus2csc ( [ doc ] , self . num_features ) . T ) <EOL> else : <EOL> doc = matutils . unitvec ( matutils . sparse2full ( doc , self . num_features ) ) <EOL> self . fresh_docs . append ( doc ) <EOL> self . fresh_nnz += doclen <EOL> if len ( self . fresh_docs ) >= self . shardsize : <EOL> self . close_shard ( ) <EOL> if len ( self . fresh_docs ) % <NUM_LIT> == <NUM_LIT:0> : <EOL> logger . info ( "<STR_LIT>" % len ( self . fresh_docs ) ) <EOL> def shardid2filename ( self , shardid ) : <EOL> if self . output_prefix . endswith ( '<STR_LIT:.>' ) : <EOL> return "<STR_LIT>" % ( self . output_prefix , shardid ) <EOL> else : <EOL> return "<STR_LIT>" % ( self . output_prefix , shardid ) <EOL> def close_shard ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . fresh_docs : <EOL> return <EOL> shardid = len ( self . shards ) <EOL> issparse = <NUM_LIT> > <NUM_LIT:1.0> * self . fresh_nnz / ( len ( self . fresh_docs ) * self . num_features ) <EOL> if issparse : <EOL> index = SparseMatrixSimilarity ( self . fresh_docs , num_terms = self . num_features , <EOL> num_docs = len ( self . fresh_docs ) , num_nnz = self . fresh_nnz ) <EOL> else : <EOL> index = MatrixSimilarity ( self . fresh_docs , num_features = self . num_features ) <EOL> logger . info ( "<STR_LIT>" % ( '<STR_LIT>' if issparse else '<STR_LIT>' , shardid ) ) <EOL> shard = Shard ( self . shardid2filename ( shardid ) , index ) <EOL> shard . num_best = self . num_best <EOL> shard . num_nnz = self . fresh_nnz <EOL> self . shards . append ( shard ) <EOL> self . fresh_docs , self . fresh_nnz = [ ] , <NUM_LIT:0> <EOL> def reopen_shard ( self ) : <EOL> assert self . shards <EOL> if self . fresh_docs : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> last_shard = self . shards [ - <NUM_LIT:1> ] <EOL> last_index = last_shard . get_index ( ) <EOL> logger . info ( "<STR_LIT>" % len ( last_shard ) ) <EOL> self . fresh_docs = list ( last_index . index ) <EOL> self . fresh_nnz = last_shard . num_nnz <EOL> del self . shards [ - <NUM_LIT:1> ] <EOL> logger . debug ( "<STR_LIT>" ) <EOL> def query_shards ( self , query ) : <EOL> """<STR_LIT>""" <EOL> args = zip ( [ query ] * len ( self . shards ) , self . shards ) <EOL> if PARALLEL_SHARDS and PARALLEL_SHARDS > <NUM_LIT:1> : <EOL> logger . debug ( "<STR_LIT>" % PARALLEL_SHARDS ) <EOL> pool = multiprocessing . Pool ( PARALLEL_SHARDS ) <EOL> result = pool . imap ( query_shard , args , chunksize = <NUM_LIT:1> + len ( args ) / PARALLEL_SHARDS ) <EOL> else : <EOL> pool = None <EOL> result = imap ( query_shard , args ) <EOL> return pool , result <EOL> def __getitem__ ( self , query ) : <EOL> """<STR_LIT>""" <EOL> self . close_shard ( ) <EOL> for shard in self . shards : <EOL> shard . num_best = self . num_best <EOL> shard . normalize = self . normalize <EOL> pool , shard_results = self . query_shards ( query ) <EOL> if self . num_best is None : <EOL> result = numpy . hstack ( shard_results ) <EOL> else : <EOL> offsets = numpy . cumsum ( [ <NUM_LIT:0> ] + [ len ( shard ) for shard in self . shards ] ) <EOL> convert = lambda doc , shard_no : [ ( doc_index + offsets [ shard_no ] , sim ) <EOL> for doc_index , sim in doc ] <EOL> is_corpus , query = utils . is_corpus ( query ) <EOL> is_corpus = is_corpus or hasattr ( query , '<STR_LIT>' ) and query . ndim > <NUM_LIT:1> and query . shape [ <NUM_LIT:0> ] > <NUM_LIT:1> <EOL> if not is_corpus : <EOL> results = ( convert ( result , shard_no ) for shard_no , result in enumerate ( shard_results ) ) <EOL> result = heapq . nlargest ( self . num_best , itertools . chain ( * results ) , key = lambda item : item [ <NUM_LIT:1> ] ) <EOL> else : <EOL> results = [ ] <EOL> for shard_no , result in enumerate ( shard_results ) : <EOL> shard_result = [ convert ( doc , shard_no ) for doc in result ] <EOL> results . append ( shard_result ) <EOL> result = [ ] <EOL> for parts in izip ( * results ) : <EOL> merged = heapq . nlargest ( self . num_best , itertools . chain ( * parts ) , key = lambda item : item [ <NUM_LIT:1> ] ) <EOL> result . append ( merged ) <EOL> if pool : <EOL> pool . terminate ( ) <EOL> return result <EOL> def vector_by_id ( self , docpos ) : <EOL> """<STR_LIT>""" <EOL> self . close_shard ( ) <EOL> pos = <NUM_LIT:0> <EOL> for shard in self . shards : <EOL> pos += len ( shard ) <EOL> if docpos < pos : <EOL> break <EOL> if not self . shards or docpos < <NUM_LIT:0> or docpos >= pos : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> ( docpos , len ( self ) ) ) <EOL> result = shard . get_document_id ( docpos - pos + len ( shard ) ) <EOL> return result <EOL> def similarity_by_id ( self , docpos ) : <EOL> """<STR_LIT>""" <EOL> query = self . vector_by_id ( docpos ) <EOL> norm , self . normalize = self . normalize , False <EOL> result = self [ query ] <EOL> self . normalize = norm <EOL> return result <EOL> def __iter__ ( self ) : <EOL> """<STR_LIT>""" <EOL> norm , self . normalize = self . normalize , False <EOL> for chunk in self . iter_chunks ( ) : <EOL> if chunk . shape [ <NUM_LIT:0> ] > <NUM_LIT:1> : <EOL> for sim in self [ chunk ] : <EOL> yield sim <EOL> else : <EOL> yield self [ chunk ] <EOL> self . normalize = norm <EOL> def iter_chunks ( self , chunksize = None ) : <EOL> """<STR_LIT>""" <EOL> self . close_shard ( ) <EOL> if chunksize is None : <EOL> chunksize = self . chunksize <EOL> for shard in self . shards : <EOL> query = shard . get_index ( ) . index <EOL> for chunk_start in xrange ( <NUM_LIT:0> , query . shape [ <NUM_LIT:0> ] , chunksize ) : <EOL> chunk_end = min ( query . shape [ <NUM_LIT:0> ] , chunk_start + chunksize ) <EOL> chunk = query [ chunk_start : chunk_end ] <EOL> yield chunk <EOL> def check_moved ( self ) : <EOL> """<STR_LIT>""" <EOL> dirname = os . path . dirname ( self . output_prefix ) <EOL> for shard in self . shards : <EOL> shard . dirname = dirname <EOL> def save ( self , fname = None , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> self . close_shard ( ) <EOL> if fname is None : <EOL> fname = self . output_prefix <EOL> super ( Similarity , self ) . save ( fname , * args , ** kwargs ) <EOL> def destroy ( self ) : <EOL> """<STR_LIT>""" <EOL> import glob <EOL> for fname in glob . glob ( self . output_prefix + '<STR_LIT:*>' ) : <EOL> logger . info ( "<STR_LIT>" % fname ) <EOL> os . remove ( fname ) <EOL> class MatrixSimilarity ( interfaces . SimilarityABC ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , corpus , num_best = None , dtype = numpy . float32 , num_features = None , chunksize = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> if num_features is None : <EOL> logger . warning ( "<STR_LIT>" ) <EOL> num_features = <NUM_LIT:1> + utils . get_max_id ( corpus ) <EOL> self . num_features = num_features <EOL> self . num_best = num_best <EOL> self . normalize = True <EOL> self . chunksize = chunksize <EOL> if corpus is not None : <EOL> if self . num_features <= <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> logger . info ( "<STR_LIT>" % <EOL> ( len ( corpus ) , num_features ) ) <EOL> self . index = numpy . empty ( shape = ( len ( corpus ) , num_features ) , dtype = dtype ) <EOL> for docno , vector in enumerate ( corpus ) : <EOL> if docno % <NUM_LIT:1000> == <NUM_LIT:0> : <EOL> logger . debug ( "<STR_LIT>" % ( docno , len ( corpus ) ) ) <EOL> if isinstance ( vector , numpy . ndarray ) : <EOL> pass <EOL> elif scipy . sparse . issparse ( vector ) : <EOL> vector = vector . toarray ( ) . flatten ( ) <EOL> else : <EOL> vector = matutils . unitvec ( matutils . sparse2full ( vector , num_features ) ) <EOL> self . index [ docno ] = vector <EOL> def __len__ ( self ) : <EOL> return self . index . shape [ <NUM_LIT:0> ] <EOL> def get_similarities ( self , query ) : <EOL> """<STR_LIT>""" <EOL> is_corpus , query = utils . is_corpus ( query ) <EOL> if is_corpus : <EOL> query = numpy . asarray ( [ matutils . sparse2full ( vec , self . num_features ) for vec in query ] , <EOL> dtype = self . index . dtype ) <EOL> else : <EOL> if scipy . sparse . issparse ( query ) : <EOL> query = query . toarray ( ) <EOL> elif isinstance ( query , numpy . ndarray ) : <EOL> pass <EOL> else : <EOL> query = matutils . sparse2full ( query , self . num_features ) <EOL> query = numpy . asarray ( query , dtype = self . index . dtype ) <EOL> result = numpy . dot ( self . index , query . T ) . T <EOL> return result <EOL> class SparseMatrixSimilarity ( interfaces . SimilarityABC ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , corpus , num_features = None , num_terms = None , num_docs = None , num_nnz = None , <EOL> num_best = None , chunksize = <NUM_LIT> , dtype = numpy . float32 ) : <EOL> self . num_best = num_best <EOL> self . normalize = True <EOL> self . chunksize = chunksize <EOL> if corpus is not None : <EOL> logger . info ( "<STR_LIT>" ) <EOL> try : <EOL> num_terms , num_docs , num_nnz = corpus . num_terms , corpus . num_docs , corpus . num_nnz <EOL> logger . debug ( "<STR_LIT>" ) <EOL> except AttributeError : <EOL> pass <EOL> if num_features is not None : <EOL> num_terms = num_features <EOL> if num_terms is None : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> corpus = ( matutils . scipy2sparse ( v ) if scipy . sparse . issparse ( v ) else <EOL> ( matutils . full2sparse ( v ) if isinstance ( v , numpy . ndarray ) else <EOL> matutils . unitvec ( v ) ) for v in corpus ) <EOL> self . index = matutils . corpus2csc ( corpus , num_terms = num_terms , num_docs = num_docs , num_nnz = num_nnz , <EOL> dtype = dtype , printprogress = <NUM_LIT> ) . T <EOL> self . index = self . index . tocsr ( ) <EOL> logger . info ( "<STR_LIT>" % self . index ) <EOL> def __len__ ( self ) : <EOL> return self . index . shape [ <NUM_LIT:0> ] <EOL> def get_similarities ( self , query ) : <EOL> """<STR_LIT>""" <EOL> is_corpus , query = utils . is_corpus ( query ) <EOL> if is_corpus : <EOL> query = matutils . corpus2csc ( query , self . index . shape [ <NUM_LIT:1> ] , dtype = self . index . dtype ) <EOL> else : <EOL> if scipy . sparse . issparse ( query ) : <EOL> query = query . T <EOL> elif isinstance ( query , numpy . ndarray ) : <EOL> if query . ndim == <NUM_LIT:1> : <EOL> query . shape = ( <NUM_LIT:1> , len ( query ) ) <EOL> query = scipy . sparse . csr_matrix ( query , dtype = self . index . dtype ) . T <EOL> else : <EOL> query = matutils . corpus2csc ( [ query ] , self . index . shape [ <NUM_LIT:1> ] , dtype = self . index . dtype ) <EOL> result = self . index * query . tocsc ( ) <EOL> if result . shape [ <NUM_LIT:1> ] == <NUM_LIT:1> and not is_corpus : <EOL> result = result . toarray ( ) . flatten ( ) <EOL> else : <EOL> result = result . toarray ( ) . T <EOL> return result </s>
