Spaces:
Runtime error
Runtime error
# | |
# 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. | |
# | |
import logging | |
from typing import List, Dict | |
from pyserini.pyclass import autoclass | |
logger = logging.getLogger(__name__) | |
JLuceneIndexer = autoclass('io.anserini.index.SimpleIndexer') | |
JsonCollectionDocument = autoclass('io.anserini.collection.JsonCollection$Document') | |
JacksonObjectMapper = autoclass('com.fasterxml.jackson.databind.ObjectMapper') | |
JacksonJsonNode = autoclass('com.fasterxml.jackson.databind.JsonNode') | |
class LuceneIndexer: | |
"""Wrapper class for ``SimpleIndexer`` in Anserini. Provides basic functionality for on-the-fly indexing via a | |
programmatic API, i.e., indexing in-process objects as opposed to on-file documents. | |
Parameters | |
---------- | |
index_dir : str | |
Path to Lucene index directory. | |
args : List[str] | |
List of arguments to pass to ``SimpleIndexer``. | |
append : bool | |
Append to existing index. | |
threads : int | |
Number of indexing threads. | |
""" | |
def __init__(self, index_dir: str = None, args: List[str] = None, append: bool = False, threads: int = 8): | |
self.index_dir = index_dir | |
self.args = args | |
if args: | |
args.extend(['-input', '', '-collection', 'JsonCollection', '-threads', str(threads)]) | |
if append: | |
args.extend(['-append']) | |
self.object = JLuceneIndexer(args) | |
else: | |
self.object = JLuceneIndexer(index_dir, append, int(threads)) | |
self.mapper = JacksonObjectMapper() | |
def add_doc_raw(self, doc: str): | |
"""Add a raw document (in the form of a JSON string) to the index. | |
Parameters | |
---------- | |
doc : str | |
Document to add. | |
""" | |
self.object.addRawDocument(doc) | |
def add_doc_dict(self, doc: Dict[str, str]): | |
"""Add a document (in the form of a Python dictionary) to the index. | |
Parameters | |
---------- | |
doc : Dict[str, str] | |
Document to add. | |
""" | |
self.object.addJsonDocument(JsonCollectionDocument.fromFields(doc['id'], doc['contents'])) | |
def add_doc_json(self, node: JacksonJsonNode): | |
"""Add a document (in the form of a Jackson JSON node object) to the index. | |
Parameters | |
---------- | |
node : JacksonJsonNode | |
Document to add. | |
""" | |
self.object.addJsonNode(node) | |
def add_batch_raw(self, docs: List[str]): | |
"""Add a batch of raw documents (in the form of JSON strings) to the index. | |
Parameters | |
---------- | |
docs : List[str] | |
Documents to add. | |
""" | |
self.object.addRawDocuments(docs) | |
def add_batch_dict(self, docs: List[Dict[str, str]]): | |
"""Add a batch of documents (in the form of Python dictionaries) to the index. | |
Parameters | |
---------- | |
docs : List[Dict[str, str]] | |
Documents to add. | |
""" | |
docs = list(map(lambda d: JsonCollectionDocument.fromFields(d['id'], d['contents']), docs)) | |
self.object.addJsonDocuments(docs) | |
def add_batch_json(self, nodes: List[JacksonJsonNode]): | |
"""Add a batch of documents (in the form of Jackson JSON node objects) to the index. | |
Parameters | |
---------- | |
nodes : List[JacksonJsonNode] | |
Documents to add. | |
""" | |
self.object.addJsonNodes(nodes) | |
def close(self): | |
"""Close this indexer, committing all in-memory data to disk.""" | |
self.object.close() | |