class ElasticSearchProcessor(MultiPackProcessor):
r"""This processor searches for relevant documents for a query"""
# pylint: disable=useless-super-delegation
def __init__(self) -> None:
super().__init__()
def initialize(self, resources: Resources, configs: Config):
self.resources = resources
self.config = configs
self.index = ElasticSearchIndexer(config=self.config.index_config)
@classmethod
def default_configs(cls) -> Dict[str, Any]:
config = super().default_configs()
config.update({
"query_pack_name": "query",
"index_config": ElasticSearchIndexer.default_configs(),
"field": "content",
"response_pack_name_prefix": "passage"
})
return config
def _process(self, input_pack: MultiPack):
r"""Searches ElasticSearch indexer to fetch documents for a query. This
query should be contained in the input multipack with name
`self.config.query_pack_name`.
This method adds new packs to `input_pack` containing the retrieved
results. Each result is added as a `ft.onto.base_ontology.Document`.
Args:
input_pack: A multipack containing query as a pack.
"""
query_pack = input_pack.get_pack(self.config.query_pack_name)
# ElasticSearchQueryCreator adds a Query entry to query pack. We now
# fetch it as the first element.
first_query: Query = query_pack.get_single(Query) # type: ignore
# pylint: disable=isinstance-second-argument-not-valid-type
# TODO: until fix: https://github.com/PyCQA/pylint/issues/3507
if not isinstance(first_query, Dict):
raise ValueError(
"The query to the elastic indexer need to be a dictionary.")
results = self.index.search(first_query.value)
hits = results["hits"]["hits"]
for idx, hit in enumerate(hits):
document = hit["_source"]
first_query.add_result(document["doc_id"], hit["_score"])
pack: DataPack = input_pack.add_pack(
f"{self.config.response_pack_name_prefix}_{idx}"
)
pack.pack_name = document["doc_id"]
content = document[self.config.field]
pack.set_text(content)
Document(pack=pack, begin=0, end=len(content))
def setUp(self):
# create indexer
file_dir_path = os.path.dirname(__file__)
data_dir = 'data_samples/ms_marco_passage_retrieval'
self.abs_data_dir = os.path.abspath(
os.path.join(file_dir_path, *([os.pardir] * 4), data_dir))
self.index_name = "final"
indexer_config = {
"batch_size": 5,
"fields": ["doc_id", "content", "pack_info"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": {
"index_name": self.index_name,
"hosts": "localhost:9200",
"algorithm": "bm25"
},
"other_kwargs": {
"request_timeout": 10,
"refresh": True
}
}
}
self.indexer = ElasticSearchIndexer(
config={"index_name": self.index_name})
nlp: Pipeline[DataPack] = Pipeline()
nlp.set_reader(MSMarcoPassageReader())
nlp.add(DataSelectorIndexProcessor(), config=indexer_config)
nlp.initialize()
self.size = 0
for _ in nlp.process_dataset(self.abs_data_dir):
self.size += 1
self.test_dir = tempfile.mkdtemp()
def default_configs(cls) -> Dict[str, Any]:
config = super().default_configs()
config.update({
"query_pack_name": "query",
"index_config": ElasticSearchIndexer.default_configs(),
"field": "content",
"response_pack_name_prefix": "passage"
})
return config
def default_configs(cls) -> Dict[str, Any]:
r"""Returns a dictionary of default hyperparameters.
.. code-block:: python
{
"batch_size": 128,
"fields": "content",
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": ElasticSearchIndexer.default_configs(),
"kwargs": {
"request_timeout": 10,
"refresh": False
}
}
}
Here:
`"batch_size"`: int
Number of examples that will be bulk added to Elasticsearch index
`"fields"`: str, list
Field name that will be used as a key while indexing the document
`"indexer"`: dict
`"name"`: str
Name of Indexer to be used.
`"hparams"`: dict
Hyperparameters to be used for the index. See
:meth:`ElasticSearchIndexer.default_hparams` for more details
`"kwargs"`: dict
Keyword arguments that will be passed to
:meth:`ElasticSearchIndexer.add_bulk` API
"""
config = super().default_configs()
config.update({
**IndexProcessor.default_configs(), "fields":
["doc_id", "content"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": ElasticSearchIndexer.default_configs(),
"other_kwargs": {
"request_timeout": 10,
"refresh": False
}
}
})
return config
def setUp(self):
file_dir_path = os.path.dirname(__file__)
data_dir = 'data_samples/ms_marco_passage_retrieval'
self.abs_data_dir = os.path.abspath(
os.path.join(file_dir_path, *([os.pardir] * 4), data_dir))
corpus_file = os.path.join(self.abs_data_dir, 'collection.tsv')
self.expected_content = set()
with open(corpus_file, 'r') as f:
for line in f.readlines():
key, value = tuple(line.split('\t', 1))
self.expected_content.add(value)
self.index_name = "test_indexer"
indexer_config = {
"batch_size": 5,
"fields": ["doc_id", "content", "pack_info"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": {
"index_name": self.index_name,
"hosts": "localhost:9200",
"algorithm": "bm25"
},
"other_kwargs": {
"request_timeout": 10,
"refresh": True
}
}
}
self.indexer = ElasticSearchIndexer(
config={"index_name": self.index_name})
self.nlp: Pipeline[DataPack] = Pipeline()
self.reader = MSMarcoPassageReader()
self.processor = DataSelectorIndexProcessor()
self.nlp.set_reader(self.reader)
self.nlp.add(self.processor, config=indexer_config)
self.nlp.initialize()
def default_config(cls) -> Dict[str, Any]:
return {
"batch_size": 10000,
"fields": ["doc_id", "content", "pack_info"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": ElasticSearchIndexer.default_configs(),
"other_kwargs": {
"request_timeout": 60,
"refresh": False
}
}
}
class ElasticSearchProcessor(MultiPackProcessor):
r"""This processor searches for relevant documents for a query"""
# pylint: disable=useless-super-delegation
def __init__(self) -> None:
super().__init__()
def initialize(self, resources: Resources, configs: HParams):
self.resources = resources
self.config = configs
self.index = ElasticSearchIndexer(hparams=self.config.index_config)
@staticmethod
def default_configs() -> Dict[str, Any]:
return {
"query_pack_name": "query",
"index_config": ElasticSearchIndexer.default_configs(),
"field": "content"
}
def _process(self, input_pack: MultiPack):
r"""Searches ElasticSearch indexer to fetch documents for a query. This
query should be contained in the input multipack with name
`self.config.query_pack_name`.
This method adds new packs to `input_pack` containing the retrieved
results. Each result is added as a `ft.onto.base_ontology.Document`.
Args:
input_pack: A multipack containing query as a pack.
"""
query_pack = input_pack.get_pack(self.config.query_pack_name)
# ElasticSearchQueryCreator adds a Query entry to query pack. We now
# fetch it as the first element.
first_query = list(query_pack.get_entries(Query))[0]
results = self.index.search(first_query.value)
hits = results["hits"]["hits"]
packs = {}
for idx, hit in enumerate(hits):
document = hit["_source"]
first_query.update_results({document["doc_id"]: hit["_score"]})
pack = DataPack(doc_id=document["doc_id"])
content = document[self.config.field]
document = Document(pack=pack, begin=0, end=len(content))
pack.add_entry(document)
pack.set_text(content)
packs[f"{self.config.response_pack_name_prefix}_{idx}"] = pack
input_pack.update_pack(packs)
class TestDataSelectorIndexProcessor(unittest.TestCase):
def setUp(self):
file_dir_path = os.path.dirname(__file__)
data_dir = 'data_samples/ms_marco_passage_retrieval'
self.abs_data_dir = os.path.abspath(
os.path.join(file_dir_path, *([os.pardir] * 4), data_dir))
corpus_file = os.path.join(self.abs_data_dir, 'collection.tsv')
self.expected_content = set()
with open(corpus_file, 'r') as f:
for line in f.readlines():
key, value = tuple(line.split('\t', 1))
self.expected_content.add(value)
self.index_name = "test_indexer"
indexer_config = {
"batch_size": 5,
"fields": ["doc_id", "content", "pack_info"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": {
"index_name": self.index_name,
"hosts": "localhost:9200",
"algorithm": "bm25"
},
"other_kwargs": {
"request_timeout": 10,
"refresh": True
}
}
}
self.indexer = ElasticSearchIndexer(
config={"index_name": self.index_name})
self.nlp: Pipeline[DataPack] = Pipeline()
self.reader = MSMarcoPassageReader()
self.processor = DataSelectorIndexProcessor()
self.nlp.set_reader(self.reader)
self.nlp.add(self.processor, config=indexer_config)
self.nlp.initialize()
def tearDown(self):
self.indexer.elasticsearch.indices.delete(index=self.index_name,
ignore=[400, 404])
def test_pipeline(self):
size = 0
for _ in self.nlp.process_dataset(self.abs_data_dir):
size += 1
retrieved_document = self.indexer.search(
query={"query": {
"match_all": {}
}},
index_name=self.index_name,
size=size)
hits = retrieved_document["hits"]["hits"]
self.assertEqual(len(hits), size)
results = set([hit["_source"]["content"] for hit in hits])
self.assertEqual(results, self.expected_content)
def initialize(self, resources: Resources, configs: Config):
self.resources = resources
self.config = configs
self.index = ElasticSearchIndexer(config=self.config.index_config)
def initialize(self, resources: Resources, configs: Config):
super().initialize(resources, configs)
self.index = ElasticSearchIndexer(config=self.configs.index_config)
def default_configs() -> Dict[str, Any]:
return {
"query_pack_name": "query",
"index_config": ElasticSearchIndexer.default_configs(),
"field": "content"
}
def initialize(self, resources: Resources, configs: HParams):
self.resources = resources
self.config = configs
self.index = ElasticSearchIndexer(hparams=self.config.index_config)
def initialize(self, resources: Resources, configs: Config):
super().initialize(resources, configs)
self.indexer = ElasticSearchIndexer(self.configs.indexer.hparams)
class ElasticSearchIndexerBase(IndexProcessor, ABC):
r"""This processor implements the basic functions to add the data packs
into an `Elasticsearch` index."""
def initialize(self, resources: Resources, configs: Config):
super().initialize(resources, configs)
self.indexer = ElasticSearchIndexer(self.configs.indexer.hparams)
@classmethod
def default_configs(cls) -> Dict[str, Any]:
r"""Returns a dictionary of default hyperparameters.
.. code-block:: python
{
"batch_size": 128,
"fields": "content",
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": ElasticSearchIndexer.default_configs(),
"kwargs": {
"request_timeout": 10,
"refresh": False
}
}
}
Here:
`"batch_size"`: int
Number of examples that will be bulk added to `Elasticsearch` index
`"fields"`: str, list
Field name that will be used as a key while indexing the document
`"indexer"`: dict
`"name"`: str
Name of Indexer to be used.
`"hparams"`: dict
Hyperparameters to be used for the index. See
:meth:`ElasticSearchIndexer.default_hparams` for more details
`"kwargs"`: dict
Keyword arguments that will be passed to
:meth:`ElasticSearchIndexer.add_bulk` API
"""
config = super().default_configs()
config.update({
"fields": ["doc_id", "content"],
"indexer": {
"name": "ElasticSearchIndexer",
"hparams": ElasticSearchIndexer.default_configs(),
"other_kwargs": {
"request_timeout": 10,
"refresh": False
}
}
})
return config
def _bulk_process(self):
self.indexer.add_bulk(self.documents,
**self.configs.indexer.other_kwargs)
def default_configs(cls) -> Dict[str, Any]:
config = super().default_configs()
config.update({
"index_config": ElasticSearchIndexer.default_configs(),
})
return config
class TestElasticSearchIndexer(unittest.TestCase):
r"""Tests Elastic Indexer."""
def setUp(self):
self.indexer = ElasticSearchIndexer(
config={"index_name": "test_index"})
def tearDown(self):
self.indexer.elasticsearch.indices.delete(
index=self.indexer.hparams.index_name, ignore=[400, 404])
def test_add(self):
document = {
"key": "This document is created to test "
"ElasticSearchIndexer"
}
self.indexer.add(document, refresh="wait_for")
retrieved_document = self.indexer.search(
query={
"query": {
"match": {
"key": "ElasticSearchIndexer"
}
},
"_source": ["key"]
})
hits = retrieved_document["hits"]["hits"]
self.assertEqual(len(hits), 1)
self.assertEqual(hits[0]["_source"], document)
def test_add_bulk(self):
size = 10000
documents = set([
f"This document {i} is created to test "
f"ElasticSearchIndexer" for i in range(size)
])
self.indexer.add_bulk([{
"key": document
} for document in documents],
refresh="wait_for")
retrieved_document = self.indexer.search(
query={"query": {
"match_all": {}
}},
index_name="test_index",
size=size)
hits = retrieved_document["hits"]["hits"]
self.assertEqual(len(hits), size)
results = set([hit["_source"]["key"] for hit in hits])
self.assertEqual(results, documents)
@performance_test
@data([100, 0.3], [500, 0.3], [1000, 0.3])
@unpack
def test_speed(self, size, epsilon):
es = Elasticsearch()
documents = [{
"_index":
"test_index_",
"_type":
"document",
"key":
f"This document {i} is created to test "
f"ElasticSearchIndexer"
} for i in range(size)]
start = time.time()
bulk(es, documents, refresh=False)
baseline = time.time() - start
es.indices.delete(index="test_index_", ignore=[400, 404])
documents = set([
f"This document {i} is created to test "
f"ElasticSearchIndexer" for i in range(size)
])
start = time.time()
self.indexer.add_bulk([{
"key": document
} for document in documents],
refresh=False)
forte_time = time.time() - start
self.assertLessEqual(forte_time, baseline + epsilon)
def setUp(self):
self.indexer = ElasticSearchIndexer(
config={"index_name": "test_index"})
def setUp(self):
self.indexer = ElasticSearchIndexer(
hparams={"index_name": "test_index"})