<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2011 <NAME> <<EMAIL>>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html

"""
This module contains functions and classes for computing similarities across
a collection of documents in the Vector Space Model.

The main class is `Similarity`, which builds an index for a given set of documents.
Once the index is built, you can perform efficient queries like "Tell me how similar
is this query document to each document in the index?". The result is a vector
of numbers as large as the size of the initial set of documents, that is, one float
for each index document. Alternatively, you can also request only the top-N most
similar index documents to the query.

You can later add new documents to the index via `Similarity.add_documents()`.

How It Works
------------

The `Similarity` class splits the index into several smaller sub-indexes ("shards"),
which are disk-based. If your entire index fits in memory (~hundreds of thousands
documents for 1GB of RAM), you can also use the `MatrixSimilarity` or `SparseMatrixSimilarity`
classes directly. These are more simple but do not scale well (the entire index is
kept in RAM).

Once the index has been initialized, you can query for document similarity simply by:

>>> index = Similarity('/tmp/tst', corpus, num_features=12) # build the index
>>> similarities = index[query] # get similarities between the query and all index documents

If you have more query documents, you can submit them all at once, in a batch:

>>> for similarities in index[batch_of_documents]: # the batch is simply an iterable of documents (=gensim corpus)
>>>     ...

The benefit of this batch (aka "chunked") querying is much better performance.
To see the speed-up on your machine, run ``python -m gensim.test.simspeed``
(compare to my results `here <http://groups.google.com/group/gensim/msg/4f6f171a869e4fca?>`_).

There is also a special syntax for when you need similarity of documents in the index
to the index itself (i.e. query=index documents themselves). This special syntax
already uses the faster, batch queries internally:

>>> for similarities in index: # return all similarities of the 1st index document, then 2nd...
>>>     ...


"""


import logging
import itertools
import os
import random
import heapq

import numpy
import scipy
import scipy.sparse

from gensim import interfaces, utils, matutils


logger = logging.getLogger('gensim.similarities.docsim')

PARALLEL_SHARDS = False
try:
    import multiprocessing
    PARALLEL_SHARDS = multiprocessing.cpu_count() # use #processes = #CPus
except ImportError:
    pass




class Shard(utils.SaveLoad):
    """A proxy class that represents a single shard instance within a Similarity
    index.

    Basically just wraps (Sparse)MatrixSimilarity so that it mmaps from disk on
    request (query).

    """
    def __init__(self, fname, index):
        self.fname = fname
        self.length = len(index)
        self.cls = index.__class__
        logger.info("saving index shard to %s" % fname)
        index.save(fname)
        # TODO: support for remote shards (Pyro? multiprocessing?)


    def __len__(self):
        return self.length


    def __str__(self):
        return ("%s Shard(%i documents in %s)" % (self.cls.__name__, len(self), self.fname))


    def get_index(self):
        logger.debug("loading index from %s" % self.fname)
        return self.cls.load(self.fname)


    def get_document_id(self, pos):
        """Return index vector at position `pos`.

        The vector is of the same type as the underlying index (ie., dense for
        MatrixSimilarity and scipy.sparse for SparseMatrixSimilarity.
        """
        assert 0 <= pos < len(self), "requested position out of range"
        index = self.get_index()
        return index.index[pos]


    def __getitem__(self, query):
        index = self.get_index()
        try:
            index.num_best = self.num_best
            index.normalize = self.normalize
        except:
            raise ValueError("num_best and normalize have to be set before querying a proxy Shard object")
        return index[query]


def query_shard(args):
    query, shard = args # simulate starmap (not part of multiprocessing in older Pythons)
    return shard[query]


class Similarity(interfaces.SimilarityABC):
    """
    Compute cosine similarity of a dynamic query against a static corpus of documents
    ("the index").

    Scalability is achieved by sharding the index into smaller pieces, each of which
    fits into core memory (see the `(Sparse)MatrixSimilarity` classes in this module).
    The shards themselves are simply stored as files to disk and mmap'ed back as needed.

    """
    def __init__(self, output_prefix, corpus, num_features, num_best=None, chunksize=1024, shardsize=32768):
        """
        Construct the index from `corpus`. The index can be later extended by calling
        the `add_documents` method. Documents are split into shards of `shardsize`
        documents each, converted to a matrix (for fast BLAS calls) and stored to disk
        under `output_prefix.shard_number` (=you need write access to that location).
        If you don't specify an output prefix, a random filename in temp will be used.

        `shardsize` should be chosen so that a `shardsize x chunksize` matrix of floats
        fits comfortably into main memory.

        `num_features` is the number of features in the `corpus` (e.g. size of the
        dictionary, or the number of latent topics for latent semantic models).

        If `num_best` is left unspecified, similarity queries will return a full
        vector with one float for every document in the index:

        >>> index = Similarity('/path/to/index', corpus, num_features=400) # if corpus has 7 documents...
        >>> index[query] # ... then result will have 7 floats
        [0.0, 0.0, 0.2, 0.13, 0.8, 0.0, 0.1]

        If `num_best` is set, queries return only the `num_best` most similar documents,
        always leaving out documents for which the similarity is 0
        (i.e. vectors with all dimensions zero and vectors with non-zero dimensions, but none in common).
        If the input vector itself has only dimensions with value zero (the sparse representation is empty),
        the returned list will always be empty.

        >>> index.num_best = 3
        >>> index[query] # return at most "num_best" of `(index_of_document, similarity)` tuples
        [(4, 0.8), (2, 0.13), (3, 0.13)]

        You can also override `num_best` dynamically, simply by setting e.g.
        `self.num_best = 10` before doing a query.

        """
        if output_prefix is None:
            # undocumented feature: set output_prefix=None to create the server in temp
            self.output_prefix = utils.randfname(prefix='simserver')
        else:
            self.output_prefix = output_prefix
        self.num_features = num_features
        self.num_best = num_best
        self.normalize = True
        self.chunksize = int(chunksize)
        self.shardsize = shardsize
        self.shards = []
        self.fresh_docs, self.fresh_nnz = [], 0

        if corpus is not None:
            self.add_documents(corpus)


    def __len__(self):
        return len(self.fresh_docs) + sum([len(shard) for shard in self.shards])


    def __str__(self):
        return ("Similarity index with %i documents in %i shards (stored under %s)" %
                (len(self), len(self.shards), self.output_prefix))


    def add_documents(self, corpus):
        """
        Extend the index with new documents.

        Internally, documents are buffered and then spilled to disk when there's
        `self.shardsize` of them (or when a query is issued).
        """
        min_ratio = 1.0 # 0.5 to only reopen shards that are <50% complete
        if self.shards and len(self.shards[-1]) < min_ratio * self.shardsize:
            # The last shard was incomplete (<; load it back and add the documents there, don't start a new shard
            self.reopen_shard()
        for doc in corpus:
            if isinstance(doc, numpy.ndarray):
                doclen = len(doc)
            elif scipy.sparse.issparse(doc):
                doclen = doc.nnz
            else:
                doclen = len(doc)
                if doclen < 0.3 * self.num_features:
                    doc = matutils.unitvec(matutils.corpus2csc([doc], self.num_features).T)
                else:
                    doc = matutils.unitvec(matutils.sparse2full(doc, self.num_features))
            self.fresh_docs.append(doc)
            self.fresh_nnz += doclen
            if len(self.fresh_docs) >= self.shardsize:
                self.close_shard()
            if len(self.fresh_docs) % 10000 == 0:
                logger.info("PROGRESS: fresh_shard size=%i" % len(self.fresh_docs))


    def shardid2filename(self, shardid):
        if self.output_prefix.endswith('.'):
            return "%s%s" % (self.output_prefix, shardid)
        else:
            return "%s.%s" % (self.output_prefix, shardid)


    def close_shard(self):
        """
        Force the latest shard to close (be converted to a matrix and stored
        to disk). Do nothing if no new documents added since last call.

        **NOTE**: the shard is closed even if it is not full yet (its size is smaller
        than `self.shardsize`). If documents are added later via `add_documents()`,
        this incomplete shard will be loaded again and completed.
        """
        if not self.fresh_docs:
            return
        shardid = len(self.shards)
        # consider the shard sparse if its density is < 30%
        issparse = 0.3 > 1.0 * self.fresh_nnz / (len(self.fresh_docs) * self.num_features)
        if issparse:
            index = SparseMatrixSimilarity(self.fresh_docs, num_terms=self.num_features,
                                           num_docs=len(self.fresh_docs), num_nnz=self.fresh_nnz)
        else:
            index = MatrixSimilarity(self.fresh_docs, num_features=self.num_features)
        logger.info("creating %s shard #%s" % ('sparse' if issparse else 'dense', shardid))
        shard = Shard(self.shardid2filename(shardid), index)
        shard.num_best = self.num_best
        shard.num_nnz = self.fresh_nnz
        self.shards.append(shard)
        self.fresh_docs, self.fresh_nnz = [], 0


    def reopen_shard(self):
        assert self.shards
        if self.fresh_docs:
            raise ValueError("cannot reopen a shard with fresh documents in index")
        last_shard = self.shards[-1]
        last_index = last_shard.get_index()
        logger.info("reopening an incomplete shard of %i documents" % len(last_shard))

        self.fresh_docs = list(last_index.index)
        self.fresh_nnz = last_shard.num_nnz
        del self.shards[-1] # remove the shard from index, *but its file on disk is not deleted*
        logger.debug("reopen complete")


    def query_shards(self, query):
        """
        Return the result of applying shard[query] for each shard in self.shards,
        as a sequence.

        If PARALLEL_SHARDS is set, the shards are queried in parallel, using
        the multiprocessing module.
        """
        args = zip([query] * len(self.shards), self.shards)
        if PARALLEL_SHARDS:
            pool = multiprocessing.Pool(PARALLEL_SHARDS)
            result = pool.map(query_shard, args) # TODO set chunksize somehow?
            # gc doesn't seem to collect the Pools... leading to "IOError 24: too many open files"
            # so let's terminate it manually. this means we cannot use imap above :(
            pool.terminate()
        else:
            # serial processing, one shard after another
            result = itertools.imap(query_shard, args)
        return result


    def __getitem__(self, query):
        """Get similarities of document `query` to all documents in the corpus.

        **or**

        If `query` is a corpus (iterable of documents), return a matrix of similarities
        of all query documents vs. all corpus document. This batch query is more
        efficient than computing the similarities one document after another.
        """
        self.close_shard() # no-op if no documents added to index since last query

        # reset num_best and normalize parameters, in case they were changed dynamically
        for shard in self.shards:
            shard.num_best = self.num_best
            shard.normalize = self.normalize

        # there are 4 distinct code paths, depending on whether input `query` is
        # a corpus (or numpy/scipy matrix) or a single document, and whether the
        # similarity result is a full array or only num_best most similar documents.

        shard_results = self.query_shards(query)
        if self.num_best is None:
            # user asked for all documents => just stack the sub-results into a single matrix
            # (works for both corpus / single doc query)
            return numpy.hstack(shard_results)

        # the following uses a lot of lazy evaluation and (optionally) parallel
        # processing, to improve query latency and minimize memory footprint.
        offsets = numpy.cumsum([0] + [len(shard) for shard in self.shards])
        convert = lambda doc, shard_no: [(doc_index + offsets[shard_no], sim)
                                         for doc_index, sim in doc]
        is_corpus, query = utils.is_corpus(query)
        is_corpus = is_corpus or hasattr(query, 'ndim') and query.ndim > 1 and query.shape[0] > 1
        if not is_corpus:
            # user asked for num_best most similar and query is a single doc
            results = (convert(result, shard_no)
                       for shard_no, result in enumerate(shard_results))
            return heapq.nlargest(self.num_best, itertools.chain(*results), key=lambda item: item[1])

        # the trickiest combination: returning num_best results when query was a corpus
        results = []
        for shard_no, result in enumerate(shard_results):
            shard_result = [convert(doc, shard_no) for doc in result]
            results.append(shard_result)
        result = []
        for parts in itertools.izip(*results):
            merged = heapq.nlargest(self.num_best, itertools.chain(*parts), key=lambda item: item[1])
            result.append(merged)
        return result


    def vector_by_id(self, docpos):
        """
        Return indexed vector corresponding to the document at position `docpos`.
        """
        self.close_shard() # no-op if no documents added to index since last query
        pos = 0
        for shard in self.shards:
            pos += len(shard)
            if docpos < pos:
                break
        if not self.shards or docpos < 0 or docpos >= pos:
            raise ValueError("invalid document position: %s (must be 0 <= x < %s)" %
                             (docpos, len(self)))
        result = shard.get_document_id(docpos - pos + len(shard))
        return result


    def similarity_by_id(self, docpos):
        """
        Return similarity of the given document only. `docpos` is the position
        of the query document within index.
        """
        query = self.vector_by_id(docpos)
        norm, self.normalize = self.normalize, False
        result = self[query]
        self.normalize = norm<commit_msg>use multiprocessing.imap instead of map<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2011 <NAME> <<EMAIL>>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html

"""
This module contains functions and classes for computing similarities across
a collection of documents in the Vector Space Model.

The main class is `Similarity`, which builds an index for a given set of documents.
Once the index is built, you can perform efficient queries like "Tell me how similar
is this query document to each document in the index?". The result is a vector
of numbers as large as the size of the initial set of documents, that is, one float
for each index document. Alternatively, you can also request only the top-N most
similar index documents to the query.

You can later add new documents to the index via `Similarity.add_documents()`.

How It Works
------------

The `Similarity` class splits the index into several smaller sub-indexes ("shards"),
which are disk-based. If your entire index fits in memory (~hundreds of thousands
documents for 1GB of RAM), you can also use the `MatrixSimilarity` or `SparseMatrixSimilarity`
classes directly. These are more simple but do not scale well (the entire index is
kept in RAM).

Once the index has been initialized, you can query for document similarity simply by:

>>> index = Similarity('/tmp/tst', corpus, num_features=12) # build the index
>>> similarities = index[query] # get similarities between the query and all index documents

If you have more query documents, you can submit them all at once, in a batch:

>>> for similarities in index[batch_of_documents]: # the batch is simply an iterable of documents (=gensim corpus)
>>>     ...

The benefit of this batch (aka "chunked") querying is much better performance.
To see the speed-up on your machine, run ``python -m gensim.test.simspeed``
(compare to my results `here <http://groups.google.com/group/gensim/msg/4f6f171a869e4fca?>`_).

There is also a special syntax for when you need similarity of documents in the index
to the index itself (i.e. query=index documents themselves). This special syntax
already uses the faster, batch queries internally:

>>> for similarities in index: # return all similarities of the 1st index document, then 2nd...
>>>     ...


"""


import logging
import itertools
import os
import random
import heapq

import numpy
import scipy
import scipy.sparse

from gensim import interfaces, utils, matutils


logger = logging.getLogger('gensim.similarities.docsim')

PARALLEL_SHARDS = False
try:
    import multiprocessing
    PARALLEL_SHARDS = multiprocessing.cpu_count() # use #processes = #CPus
except ImportError:
    pass




class Shard(utils.SaveLoad):
    """A proxy class that represents a single shard instance within a Similarity
    index.

    Basically just wraps (Sparse)MatrixSimilarity so that it mmaps from disk on
    request (query).

    """
    def __init__(self, fname, index):
        self.fname = fname
        self.length = len(index)
        self.cls = index.__class__
        logger.info("saving index shard to %s" % fname)
        index.save(fname)
        # TODO: support for remote shards (Pyro? multiprocessing?)


    def __len__(self):
        return self.length


    def __str__(self):
        return ("%s Shard(%i documents in %s)" % (self.cls.__name__, len(self), self.fname))


    def get_index(self):
        logger.debug("loading index from %s" % self.fname)
        return self.cls.load(self.fname)


    def get_document_id(self, pos):
        """Return index vector at position `pos`.

        The vector is of the same type as the underlying index (ie., dense for
        MatrixSimilarity and scipy.sparse for SparseMatrixSimilarity.
        """
        assert 0 <= pos < len(self), "requested position out of range"
        index = self.get_index()
        return index.index[pos]


    def __getitem__(self, query):
        index = self.get_index()
        try:
            index.num_best = self.num_best
            index.normalize = self.normalize
        except:
            raise ValueError("num_best and normalize have to be set before querying a proxy Shard object")
        return index[query]


def query_shard(args):
    query, shard = args # simulate starmap (not part of multiprocessing in older Pythons)
    logger.debug("querying shard %s num_best=%s in process %s" % (shard, shard.num_best, os.getpid()))
    result = shard[query]
    logger.debug("finished querying shard %s in process %s" % (shard, os.getpid()))
    return result



class Similarity(interfaces.SimilarityABC):
    """
    Compute cosine similarity of a dynamic query against a static corpus of documents
    ("the index").

    Scalability is achieved by sharding the index into smaller pieces, each of which
    fits into core memory (see the `(Sparse)MatrixSimilarity` classes in this module).
    The shards themselves are simply stored as files to disk and mmap'ed back as needed.

    """
    def __init__(self, output_prefix, corpus, num_features, num_best=None, chunksize=1024, shardsize=32768):
        """
        Construct the index from `corpus`. The index can be later extended by calling
        the `add_documents` method. Documents are split into shards of `shardsize`
        documents each, converted to a matrix (for fast BLAS calls) and stored to disk
        under `output_prefix.shard_number` (=you need write access to that location).
        If you don't specify an output prefix, a random filename in temp will be used.

        `shardsize` should be chosen so that a `shardsize x chunksize` matrix of floats
        fits comfortably into main memory.

        `num_features` is the number of features in the `corpus` (e.g. size of the
        dictionary, or the number of latent topics for latent semantic models).

        If `num_best` is left unspecified, similarity queries will return a full
        vector with one float for every document in the index:

        >>> index = Similarity('/path/to/index', corpus, num_features=400) # if corpus has 7 documents...
        >>> index[query] # ... then result will have 7 floats
        [0.0, 0.0, 0.2, 0.13, 0.8, 0.0, 0.1]

        If `num_best` is set, queries return only the `num_best` most similar documents,
        always leaving out documents for which the similarity is 0
        (i.e. vectors with all dimensions zero and vectors with non-zero dimensions, but none in common).
        If the input vector itself has only dimensions with value zero (the sparse representation is empty),
        the returned list will always be empty.

        >>> index.num_best = 3
        >>> index[query] # return at most "num_best" of `(index_of_document, similarity)` tuples
        [(4, 0.8), (2, 0.13), (3, 0.13)]

        You can also override `num_best` dynamically, simply by setting e.g.
        `self.num_best = 10` before doing a query.

        """
        if output_prefix is None:
            # undocumented feature: set output_prefix=None to create the server in temp
            self.output_prefix = utils.randfname(prefix='simserver')
        else:
            self.output_prefix = output_prefix
        self.num_features = num_features
        self.num_best = num_best
        self.normalize = True
        self.chunksize = int(chunksize)
        self.shardsize = shardsize
        self.shards = []
        self.fresh_docs, self.fresh_nnz = [], 0

        if corpus is not None:
            self.add_documents(corpus)


    def __len__(self):
        return len(self.fresh_docs) + sum([len(shard) for shard in self.shards])


    def __str__(self):
        return ("Similarity index with %i documents in %i shards (stored under %s)" %
                (len(self), len(self.shards), self.output_prefix))


    def add_documents(self, corpus):
        """
        Extend the index with new documents.

        Internally, documents are buffered and then spilled to disk when there's
        `self.shardsize` of them (or when a query is issued).
        """
        min_ratio = 1.0 # 0.5 to only reopen shards that are <50% complete
        if self.shards and len(self.shards[-1]) < min_ratio * self.shardsize:
            # The last shard was incomplete (<; load it back and add the documents there, don't start a new shard
            self.reopen_shard()
        for doc in corpus:
            if isinstance(doc, numpy.ndarray):
                doclen = len(doc)
            elif scipy.sparse.issparse(doc):
                doclen = doc.nnz
            else:
                doclen = len(doc)
                if doclen < 0.3 * self.num_features:
                    doc = matutils.unitvec(matutils.corpus2csc([doc], self.num_features).T)
                else:
                    doc = matutils.unitvec(matutils.sparse2full(doc, self.num_features))
            self.fresh_docs.append(doc)
            self.fresh_nnz += doclen
            if len(self.fresh_docs) >= self.shardsize:
                self.close_shard()
            if len(self.fresh_docs) % 10000 == 0:
                logger.info("PROGRESS: fresh_shard size=%i" % len(self.fresh_docs))


    def shardid2filename(self, shardid):
        if self.output_prefix.endswith('.'):
            return "%s%s" % (self.output_prefix, shardid)
        else:
            return "%s.%s" % (self.output_prefix, shardid)


    def close_shard(self):
        """
        Force the latest shard to close (be converted to a matrix and stored
        to disk). Do nothing if no new documents added since last call.

        **NOTE**: the shard is closed even if it is not full yet (its size is smaller
        than `self.shardsize`). If documents are added later via `add_documents()`,
        this incomplete shard will be loaded again and completed.
        """
        if not self.fresh_docs:
            return
        shardid = len(self.shards)
        # consider the shard sparse if its density is < 30%
        issparse = 0.3 > 1.0 * self.fresh_nnz / (len(self.fresh_docs) * self.num_features)
        if issparse:
            index = SparseMatrixSimilarity(self.fresh_docs, num_terms=self.num_features,
                                           num_docs=len(self.fresh_docs), num_nnz=self.fresh_nnz)
        else:
            index = MatrixSimilarity(self.fresh_docs, num_features=self.num_features)
        logger.info("creating %s shard #%s" % ('sparse' if issparse else 'dense', shardid))
        shard = Shard(self.shardid2filename(shardid), index)
        shard.num_best = self.num_best
        shard.num_nnz = self.fresh_nnz
        self.shards.append(shard)
        self.fresh_docs, self.fresh_nnz = [], 0


    def reopen_shard(self):
        assert self.shards
        if self.fresh_docs:
            raise ValueError("cannot reopen a shard with fresh documents in index")
        last_shard = self.shards[-1]
        last_index = last_shard.get_index()
        logger.info("reopening an incomplete shard of %i documents" % len(last_shard))

        self.fresh_docs = list(last_index.index)
        self.fresh_nnz = last_shard.num_nnz
        del self.shards[-1] # remove the shard from index, *but its file on disk is not deleted*
        logger.debug("reopen complete")


    def query_shards(self, query):
        """
        Return the result of applying shard[query] for each shard in self.shards,
        as a sequence.

        If PARALLEL_SHARDS is set, the shards are queried in parallel, using
        the multiprocessing module.
        """
        args = zip([query] * len(self.shards), self.shards)
        if PARALLEL_SHARDS:
            logger.debug("spawning %i query processes" % PARALLEL_SHARDS)
            pool = multiprocessing.Pool(PARALLEL_SHARDS)
            result = pool.imap(query_shard, args, chunksize=1 + len(args) / PARALLEL_SHARDS)
        else:
            # serial processing, one shard after another
            pool = None
            result = itertools.imap(query_shard, args)
        return pool, result


    def __getitem__(self, query):
        """Get similarities of document `query` to all documents in the corpus.

        **or**

        If `query` is a corpus (iterable of documents), return a matrix of similarities
        of all query documents vs. all corpus document. This batch query is more
        efficient than computing the similarities one document after another.
        """
        self.close_shard() # no-op if no documents added to index since last query

        # reset num_best and normalize parameters, in case they were changed dynamically
        for shard in self.shards:
            shard.num_best = self.num_best
            shard.normalize = self.normalize

        # there are 4 distinct code paths, depending on whether input `query` is
        # a corpus (or numpy/scipy matrix) or a single document, and whether the
        # similarity result should be a full array or only num_best most similar
        # documents.
        pool, shard_results = self.query_shards(query)
        if self.num_best is None:
            # user asked for all documents => just stack the sub-results into a single matrix
            # (works for both corpus / single doc query)
            result = numpy.hstack(shard_results)
        else:
            # the following uses a lot of lazy evaluation and (optionally) parallel
            # processing, to improve query latency and minimize memory footprint.
            offsets = numpy.cumsum([0] + [len(shard) for shard in self.shards])
            convert = lambda doc, shard_no: [(doc_index + offsets[shard_no], sim)
                                             for doc_index, sim in doc]
            is_corpus, query = utils.is_corpus(query)
            is_corpus = is_corpus or hasattr(query, 'ndim') and query.ndim > 1 and query.shape[0] > 1
            if not is_corpus:
                # user asked for num_best most similar and query is a single doc
                results = (convert(result, shard_no) for shard_no, result in enumerate(shard_results))
                result = heapq.nlargest(self.num_best, itertools.chain(*results), key=lambda item: item[1])
            else:
                # the trickiest combination: returning num_best results when query was a corpus
                results = []
                for shard_no, result in enumerate(shard_results):
                    shard_result = [convert(doc, shard_no) for doc in result]
                    results.append(shard_result)
                result = []
                for parts in itertools.izip(*results):
                    merged = heapq.nlargest(self.num_best, itertools.chain(*parts), key=lambda item: item[1])
                    result.append(merged)
        if pool:
            # gc doesn't seem to collect the Pools, eventually leading to
            # "IOError 24: too many open files". so let's terminate it manually.
            pool.terminate()
        return result


    def vector_by_id(self, docpos):
        """
        Return indexed vector corresponding to the document at position `docpos`.
        """
        self.close_shard() # no-op if no documents added to index since last query
        pos = 0
        for shard in self.shards:
            pos += len(shard)
            if docpos < pos:
                break
        if not self.shards or docpos < 0 or docpos >= pos:
            raise ValueError("invalid document position: %s (must be 0 <= x < %s)" %
                             (docpos, len(self)))
        result = shard.get_document_id(docpos - pos + len(shard))
        return result


    def similarity_by_id(self, docpos):
        """
        Return similarity of the given document only. `docpos` is the position
        of the query document within index.
        """
        query = self.vector_by_id(docpos)
        norm, self.normalize = self.normalize, False
        result = self[query]
        self.normalize = norm
{'fraction_non_alphanumeric': 0.08442118030066069, 'fraction_numerical': 0.006862213143531965, 'mean_word_length': 2.48753339269813, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 83, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2987036', 'n_tokens_mistral': 9176, 'n_tokens_neox': 8636, 'n_words': 3268}