add es retriever

.env.example

@@ -1,4 +1,5 @@
 ELASTIC_USERNAME=elastic
 ELASTIC_PASSWORD=<password>
+ELASTIC_HOST=<localhost>
 
 LOG_LEVEL=INFO

main.py

@@ -1,18 +1,18 @@
+import os
+import random
+from typing import cast
+
+import torch
+import transformers
 from datasets import DatasetDict, load_dataset
 
+from src.evaluation import evaluate
 from src.readers.dpr_reader import DprReader
+from src.retrievers.es_retriever import ESRetriever
 from src.retrievers.faiss_retriever import FaissRetriever
 from src.utils.log import get_logger
-from src.evaluation import evaluate
-from typing import cast
-
 from src.utils.preprocessing import result_to_reader_input
 
-import torch
-import transformers
-import os
-import random
-
 os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = '1'
 
 logger = get_logger()
@@ -31,7 +31,8 @@ if __name__ == '__main__':
         "GroNLP/ik-nlp-22_slp", "paragraphs"))
 
     # Initialize retriever
-    retriever = FaissRetriever(dataset_paragraphs)
+    # retriever = FaissRetriever(dataset_paragraphs)
+    retriever = ESRetriever(dataset_paragraphs)
 
     # Retrieve example
     # random.seed(111)

src/retrievers/es_retriever.py

@@ -1,15 +1,35 @@
-from datasets import load_dataset
+from datasets import DatasetDict, load_dataset
 from src.utils.log import get_logger
 from src.retrievers.base_retriever import Retriever
+from elasticsearch import Elasticsearch
+from dotenv import load_dotenv
+import os
 
+load_dotenv()
 
 logger = get_logger()
 
 
 class ESRetriever(Retriever):
-    def __init__(self, data_set: ) -> None:
+    def __init__(self, dataset: DatasetDict) -> None:
+        self.dataset = dataset["train"]
 
-        pass
+        es_host = os.getenv("ELASTIC_HOST", "localhost")
+        es_password = os.getenv("ELASTIC_PASSWORD")
+        es_username = os.getenv("ELASTIC_USERNAME")
 
-    def retrieve(self, query: str, k: int):
-        pass
+        self.client = Elasticsearch(
+            hosts=[es_host], http_auth=(es_username, es_password))
+
+        if self.client.indices.exists(index="paragraphs"):
+            self.dataset.load_elasticsearch_index(
+                "paragraphs", es_index_name="paragraphs", es_client=self.client)
+        else:
+            logger.info(f"Creating index 'paragraphs' on {es_host}")
+            self.dataset.add_elasticsearch_index(column="text",
+                                                 index_name="paragraphs",
+                                                 es_index_name="paragraphs",
+                                                 es_client=self.client)
+
+    def retrieve(self, query: str, k: int = 5):
+        return self.dataset.get_nearest_examples("paragraphs", query, k)