ArthurChen189's picture
upload pyserini
62977bb
#
# 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()