pyserini/index/lucene/_base.py

#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
This module provides Pyserini's Python interface for raw access to Lucene indexes built by Anserini. The main entry
point is the ``IndexReaderUtils`` class, which wraps the Java class with the same name in Anserini. Many of the classes
and methods provided are meant only to provide tools for examining an index and are not optimized for computing over.
"""

import logging
from enum import Enum
from typing import Dict, Iterator, List, Optional, Tuple
from tqdm import tqdm
import json
import math

from pyserini.analysis import get_lucene_analyzer, JAnalyzer, JAnalyzerUtils
from pyserini.pyclass import autoclass
from pyserini.util import download_prebuilt_index, get_sparse_indexes_info
from pyserini.prebuilt_index_info import TF_INDEX_INFO, IMPACT_INDEX_INFO

logger = logging.getLogger(__name__)


# Wrappers around Anserini classes
JDocument = autoclass('org.apache.lucene.document.Document')
JIndexReader = autoclass('io.anserini.index.IndexReaderUtils')


class JIndexHelpers:
    @staticmethod
    def JArgs():
        args = autoclass('io.anserini.index.IndexCollection$Args')()
        args.storeContents = True
        args.storeRaw = True
        args.dryRun = True ## So that indexing will be skipped

        return args

    @staticmethod
    def JCounters():
        IndexCollection = autoclass('io.anserini.index.IndexCollection')
        Counters = autoclass('io.anserini.index.IndexCollection$Counters')

        return Counters(IndexCollection)


class Document:
    """Wrapper class for a Lucene ``Document``.

    Parameters
    ----------
    document : JDocument
        Underlying Lucene ``Document``.
    """

    def __init__(self, document):
        if document is None:
            raise ValueError('Cannot create a Document with None.')
        self.object = document

    def docid(self: JDocument) -> str:
        return self.object.getField('id').stringValue()

    def id(self: JDocument) -> str:
        # Convenient alias for docid()
        return self.object.getField('id').stringValue()

    def lucene_document(self: JDocument) -> JDocument:
        return self.object

    def contents(self: JDocument) -> str:
        return self.object.get('contents')

    def raw(self: JDocument) -> str:
        return self.object.get('raw')

    def get(self: JDocument, field: str) -> str:
        return self.object.get(field)


class JGenerators(Enum):
    AclAnthologyGenerator = autoclass('io.anserini.index.generator.AclAnthologyGenerator')
    DefaultLuceneDocumentGenerator = autoclass('io.anserini.index.generator.DefaultLuceneDocumentGenerator')
    TweetGenerator = autoclass('io.anserini.index.generator.TweetGenerator')
    WashingtonPostGenerator = autoclass('io.anserini.index.generator.WashingtonPostGenerator')


class Generator:
    """Wrapper class for Anserini's generators.

    Parameters
    ----------
    generator_class : str
        Name of generator class to instantiate
    """

    def __init__(self, generator_class):
        self.counters = JIndexHelpers.JCounters()
        self.args = JIndexHelpers.JArgs()
        self.generator_class = generator_class
        self.object = self._get_generator()

    def _get_generator(self):
        try:
            return JGenerators[self.generator_class].value(self.args)
        except:
            raise ValueError(self.generator_class)

    def create_document(self, document):
        """
        Parameters
        ----------
        document : pyserini.collection.pycollection.Document
            Collection document to create Lucene document from

        Returns
        -------
        result : org.apache.lucene.document.Document
            Lucene document generated
        """
        return self.object.createDocument(document.object)


class IndexTerm:
    """Class representing an analyzed term in an index with associated statistics.

    Parameters
    ----------
    term : str
        Analyzed term.
    df : int
        Document frequency, the number of documents in the collection that contains the term.
    cf : int
        Collection frequency, the number of times the term occurs in the entire collection.  This value is equal to the
        sum of all the term frequencies of the term across all documents in the collection.
    """

    def __init__(self, term, df, cf):
        self.term = term
        self.df = df
        self.cf = cf


class Posting:
    """Class representing a posting in a postings list.

    Parameters
    ----------
    docid : int
        Collection ``docid``.
    tf : int
        Term frequency.
    positions : List[int]
        List of positions.
    """

    def __init__(self, docid, tf, positions):
        self.docid = docid
        self.tf = tf
        self.positions = positions

    def __repr__(self):
        repr = '(' + str(self.docid) + ', ' + str(self.tf) + ')'
        if self.positions:
            repr += ' [' + ','.join([str(p) for p in self.positions]) + ']'
        return repr


class IndexReader:
    """Wrapper class for ``IndexReaderUtils`` in Anserini.

    Parameters
    ----------
    index_dir : str
        Path to Lucene index directory.
    """

    def __init__(self, index_dir):
        self.object = JIndexReader()
        self.reader = self.object.getReader(index_dir)

    @classmethod
    def from_prebuilt_index(cls, prebuilt_index_name: str, verbose=False):
        """Build an index reader from a prebuilt index; download the index if necessary.

        Parameters
        ----------
        prebuilt_index_name : str
            Prebuilt index name.
        verbose : bool
            Print status information.

        Returns
        -------
        IndexReader
            Index reader built from the prebuilt index.
        """
        if verbose:
            print(f'Attempting to initialize pre-built index {prebuilt_index_name}.')

        try:
            index_dir = download_prebuilt_index(prebuilt_index_name, verbose=verbose)
        except ValueError as e:
            print(str(e))
            return None

        if verbose:
            print(f'Initializing {prebuilt_index_name}...')

        index_reader = cls(index_dir)
        # Validate index stats; will throw exception there are any issues.
        index_reader.validate(prebuilt_index_name, verbose=verbose)

        return index_reader

    @staticmethod
    def list_prebuilt_indexes():
        """Display information about available prebuilt indexes."""
        get_sparse_indexes_info()

    def analyze(self, text: str, analyzer=None) -> List[str]:
        """Analyze a piece of text. Applies Anserini's default Lucene analyzer if analyzer not specified.

        Parameters
        ----------
        text : str
            Text to analyze.
        analyzer : analyzer
            Analyzer to apply.
        Returns
        -------
        List[str]
            List of tokens corresponding to the output of the analyzer.
        """
        if analyzer is None:
            results = JAnalyzerUtils.analyze(text)
        else:
            results = JAnalyzerUtils.analyze(analyzer, text)
        tokens = []
        for token in results.toArray():
            tokens.append(token)
        return tokens

    def validate(self, prebuilt_index_name: str, verbose=False):
        """Validate this index against stored stats for a pre-built index."""
        stats = self.stats()

        if prebuilt_index_name in TF_INDEX_INFO:
            if stats['documents'] != TF_INDEX_INFO[prebuilt_index_name]['documents']:
                raise ValueError('Pre-built index fails consistency check: "documents" does not match!')
            if stats['unique_terms'] != TF_INDEX_INFO[prebuilt_index_name]['unique_terms']:
                raise ValueError('Pre-built index fails consistency check: "unique_terms" does not match!')
            if stats['total_terms'] != TF_INDEX_INFO[prebuilt_index_name]['total_terms']:
                raise ValueError('Pre-built index fails consistency check: "total_terms" does not match!')
        elif prebuilt_index_name in IMPACT_INDEX_INFO:
            if stats['documents'] != IMPACT_INDEX_INFO[prebuilt_index_name]['documents']:
                raise ValueError('Pre-built index fails consistency check: "documents" does not match!')
            if stats['unique_terms'] != IMPACT_INDEX_INFO[prebuilt_index_name]['unique_terms']:
                raise ValueError('Pre-built index fails consistency check: "unique_terms" does not match!')
            if stats['total_terms'] != IMPACT_INDEX_INFO[prebuilt_index_name]['total_terms']:
                raise ValueError('Pre-built index fails consistency check: "total_terms" does not match!')
        else:
            print(f'Unknown pre-built index \'{prebuilt_index_name}\'!')
            return False

        if verbose:
            print(stats)
            print(f'Index passes consistency checks against pre-built index \'{prebuilt_index_name}\'!')

        return True

    def terms(self) -> Iterator[IndexTerm]:
        """Return an iterator over analyzed terms in the index.

        Returns
        -------
        Iterator[IndexTerm]
            Iterator over :class:`IndexTerm` objects corresponding to (analyzed) terms in the index.
        """
        term_iterator = self.object.getTerms(self.reader)
        while term_iterator.hasNext():
            cur_term = term_iterator.next()
            yield IndexTerm(cur_term.getTerm(), cur_term.getDF(), cur_term.getTotalTF())

    def get_term_counts(self, term: str, analyzer: Optional[JAnalyzer] = get_lucene_analyzer()) -> Tuple[int, int]:
        """Return the document frequency and collection frequency of a term. Applies Anserini's default Lucene
        ``Analyzer`` if analyzer is not specified.

        Parameters
        ----------
        term : str
            Unanalyzed term.
        analyzer : analyzer
            Analyzer to apply.

        Returns
        -------
        Tuple[int, int]
            Document frequency and collection frequency.
        """
        if analyzer is None:
            analyzer = get_lucene_analyzer(stemming=False, stopwords=False)

        term_map = self.object.getTermCountsWithAnalyzer(self.reader, term, analyzer)

        return term_map.get('docFreq'), term_map.get('collectionFreq')

    def get_postings_list(self, term: str, analyzer=get_lucene_analyzer()) -> List[Posting]:
        """Return the postings list for a term.

        Parameters
        ----------
        term : str
            Raw term.
        analyzer : analyzer
            Analyzer to apply. Defaults to Anserini's default.

        Returns
        -------
        List[Posting]
            List of :class:`Posting` objects corresponding to the postings list for the term.
        """
        if analyzer is None:
            postings_list = self.object.getPostingsListForAnalyzedTerm(self.reader, term)
        else:
            postings_list = self.object.getPostingsListWithAnalyzer(self.reader, term,
                                                                    analyzer)

        if postings_list is None:
            return None

        result = []
        for posting in postings_list.toArray():
            result.append(Posting(posting.getDocid(), posting.getTF(), posting.getPositions()))
        return result

    def get_document_vector(self, docid: str) -> Optional[Dict[str, int]]:
        """Return the document vector for a ``docid``. Note that requesting the document vector of a ``docid`` that
        does not exist in the index will return ``None`` (as opposed to an empty dictionary); this forces the caller
        to handle ``None`` explicitly and guards against silent errors.

        Parameters
        ----------
        docid : str
            Collection ``docid``.

        Returns
        -------
        Optional[Dict[str, int]]
            A dictionary with analyzed terms as keys and their term frequencies as values.
        """
        doc_vector_map = self.object.getDocumentVector(self.reader, docid)
        if doc_vector_map is None:
            return None
        doc_vector_dict = {}
        for term in doc_vector_map.keySet().toArray():
            doc_vector_dict[term] = doc_vector_map.get(term)
        return doc_vector_dict

    def get_term_positions(self, docid: str) -> Optional[Dict[str, int]]:
        """Return the term position mapping of the document with ``docid``. Note that the term in the document is
        stemmed and stop words may be removed according to your index settings. Also, requesting the document vector of
        a ``docid`` that does not exist in the index will return ``None`` (as opposed to an empty dictionary); this
        forces the caller to handle ``None`` explicitly and guards against silent errors.

        Parameters
        ----------
        docid : str
            Collection ``docid``.

        Returns
        -------
        Optional[Dict[str, int]]
            A tuple contains a dictionary with analyzed terms as keys and corresponding posting list as values
        """
        java_term_position_map = self.object.getTermPositions(self.reader, docid)
        if java_term_position_map is None:
            return None
        term_position_map = {}
        for term in java_term_position_map.keySet().toArray():
            term_position_map[term] = java_term_position_map.get(term).toArray()
        return term_position_map

    def doc(self, docid: str) -> Optional[Document]:
        """Return the :class:`Document` corresponding to ``docid``. Returns ``None`` if the ``docid`` does not exist
        in the index.

        Parameters
        ----------
        docid : str
            The collection ``docid``.

        Returns
        -------
        Optional[Document]
            :class:`Document` corresponding to the ``docid``.
        """
        lucene_document = self.object.document(self.reader, docid)
        if lucene_document is None:
            return None
        return Document(lucene_document)

    def doc_by_field(self, field: str, q: str) -> Optional[Document]:
        """Return the :class:`Document` based on a ``field`` with ``id``. For example, this method can be used to fetch
        document based on alternative primary keys that have been indexed, such as an article's DOI.

        Parameters
        ----------
        field : str
            The field to look up.
        q : str
            The document's unique id.

        Returns
        -------
        Optional[Document]
            :class:`Document` whose ``field`` is ``id``.
        """
        lucene_document = self.object.documentByField(self.reader, field, q)
        if lucene_document is None:
            return None
        return Document(lucene_document)

    def doc_raw(self, docid: str) -> Optional[str]:
        """Return the raw document contents for a collection ``docid``.

        Parameters
        ----------
        docid : str
            Collection ``docid``.

        Returns
        -------
        Optional[str]
            Raw document contents.
        """
        return self.object.documentRaw(self.reader, docid)

    def doc_contents(self, docid: str) -> Optional[str]:
        """Return the indexed document contents for a collection ``docid``.

        Parameters
        ----------
        docid : str
            The collection ``docid``.

        Returns
        -------
        Optional[str]
            Index document contents.
        """
        return self.object.documentContents(self.reader, docid)

    def compute_bm25_term_weight(self, docid: str, term: str, analyzer=get_lucene_analyzer(), k1=0.9, b=0.4) -> float:
        """Compute the BM25 weight of a term in a document. Specify ``analyzer=None`` for an already analyzed term,
        e.g., from the output of :func:`get_document_vector`.

        Parameters
        ----------
        docid : str
            Collection ``docid``.
        term : str
            Term.
        analyzer : analyzer
            Lucene analyzer to use, ``None`` if term is already analyzed.
        k1 : float
            BM25 k1 parameter.
        b : float
            BM25 b parameter.

        Returns
        -------
        float
            BM25 weight of the term in the document, or 0 if the term does not exist in the document.
        """
        if analyzer is None:
            return self.object.getBM25AnalyzedTermWeightWithParameters(self.reader, docid,
                                                                       term,
                                                                       float(k1), float(b))
        else:
            return self.object.getBM25UnanalyzedTermWeightWithParameters(self.reader, docid,
                                                                         term, analyzer,
                                                                         float(k1), float(b))

    def compute_query_document_score(self, docid: str, query: str, similarity=None):
        if similarity is None:
            return self.object.computeQueryDocumentScore(self.reader, docid, query)
        else:
            return self.object.computeQueryDocumentScoreWithSimilarity(self.reader, docid, query, similarity)

    def convert_internal_docid_to_collection_docid(self, docid: int) -> str:
        """Convert Lucene's internal ``docid`` to its external collection ``docid``.

        Parameters
        ----------
        docid : int
            Lucene internal ``docid``.

        Returns
        -------
        str
            External collection ``docid`` corresponding to Lucene's internal ``docid``.
        """
        return self.object.convertLuceneDocidToDocid(self.reader, docid)

    def convert_collection_docid_to_internal_docid(self, docid: str) -> int:
        """Convert external collection ``docid`` to its Lucene's internal ``docid``.

        Parameters
        ----------
        docid : str
            External collection ``docid``.

        Returns
        -------
        str
            Lucene internal ``docid`` corresponding to the external collection ``docid``.
        """
        return self.object.convertDocidToLuceneDocid(self.reader, docid)

    def stats(self) -> Dict[str, int]:
        """Return dictionary with index statistics.

        Returns
        -------
        Dict[str, int]
            Index statistics as a dictionary of statistic's name to statistic.
            - documents: number of documents
            - non_empty_documents: number of non-empty documents
            - unique_terms: number of unique terms
            - total_terms: number of total terms
        """
        index_stats_map = self.object.getIndexStats(self.reader)

        if index_stats_map is None:
            return None

        index_stats_dict = {}
        for term in index_stats_map.keySet().toArray():
            index_stats_dict[term] = index_stats_map.get(term)

        return index_stats_dict

    def dump_documents_BM25(self, file_path, k1=0.9, b=0.4):
        """Dumps out all the document vectors with BM25 weights in Pyserini's JSONL vector format.

        Parameters
        ----------
        file_path : str
            File path to dump JSONL file.
        k1 : float
            BM25 k1 parameter.
        b : float
            BM25 b parameter.
        """

        f = open(file_path, 'w')

        assert 'documents' in self.stats()
        for i in tqdm(range(self.stats()['documents'])):
            docid = self.convert_internal_docid_to_collection_docid(i)
            bm25_vector = {}
            for term in self.get_document_vector(docid):
                bm25_vector[term] = self.compute_bm25_term_weight(docid, term, analyzer=None, k1=k1, b=b)

            # vectors are written line by line to avoid running out of memory
            f.write(json.dumps({'id': docid, 'vector': bm25_vector}) + "\n")

        f.close()

    def quantize_weights(self, input_file_path, output_file_path, bits = 8):
        """Takes vectors of weights in Pyserini's JSONL vector format and quantizes them.

        Parameters
        ----------
        input_file_path : str
            File path of vectors of weights in Pyserini's JSONL vector format.
        output_file_path : str
            File path to output JSONL file of quantized weight vectors.
        bits : int
            Number of bits to use to represent quantized scores.
        """

        min_weight = float('inf')
        max_weight = float('-inf')

        input_file = open(input_file_path, 'r')

        # vectors are read line by line to avoid running out of memory
        for line in input_file:
            doc = json.loads(line)
            for weight in doc['vector'].values():
                if weight > max_weight:
                    max_weight = weight
                if weight < min_weight:
                    min_weight = weight
        input_file.seek(0)

        output_file = open(output_file_path, 'w')

        smallest_impact = 1
        for line in input_file:
            doc = json.loads(line)
            for element in doc['vector']:
                doc['vector'][element] = math.floor((2 ** bits - smallest_impact) * (doc['vector'][element] - min_weight) / (max_weight - min_weight)) + smallest_impact
            output_file.write(json.dumps(doc) + "\n")

        input_file.close()
        output_file.close()