Fix timings and add timing results

.env.example

@@ -3,6 +3,6 @@ ELASTIC_PASSWORD=<password>
 ELASTIC_HOST=https://localhost:9200
 
 LOG_LEVEL=INFO
-TRANSFORMERS_NO_ADVISORY_WARNINGS
-
+TRANSFORMERS_NO_ADVISORY_WARNINGS=true
+KMP_DUPLICATE_LIB_OK=true
 ENABLE_TIMING=TRUE

main.py

@@ -1,25 +1,34 @@
-import random
-from typing import Dict, cast
+from dotenv import load_dotenv
+# needs to happen as very first thing, otherwise HF ignores env vars
+load_dotenv()
+
+import os
+import pandas as pd
 
-import torch
-import transformers
+from dataclasses import dataclass
+from typing import Dict, cast
 from datasets import DatasetDict, load_dataset
-from dotenv import load_dotenv
-from query import print_answers
 
+from src.readers.base_reader import Reader
 from src.evaluation import evaluate
 from src.readers.dpr_reader import DprReader
+from src.readers.longformer_reader import LongformerReader
 from src.retrievers.base_retriever import Retriever
 from src.retrievers.es_retriever import ESRetriever
-from src.retrievers.faiss_retriever import FaissRetriever
-from src.utils.log import get_logger
+from src.retrievers.faiss_retriever import (
+    FaissRetriever,
+    FaissRetrieverOptions
+)
+from src.utils.log import logger
 from src.utils.preprocessing import context_to_reader_input
 from src.utils.timing import get_times, timeit
 
-logger = get_logger()
 
-load_dotenv()
-transformers.logging.set_verbosity_error()
+@dataclass
+class Experiment:
+    retriever: Retriever
+    reader: Reader
+
 
 if __name__ == '__main__':
     dataset_name = "GroNLP/ik-nlp-22_slp"
@@ -28,41 +37,72 @@ if __name__ == '__main__':
     questions = cast(DatasetDict, load_dataset(dataset_name, "questions"))
 
     # Only doing a few questions for speed
-    subset_idx = 3
+    subset_idx = len(questions["test"])
     questions_test = questions["test"][:subset_idx]
 
-    experiments: Dict[str, Retriever] = {
-        "faiss": FaissRetriever(paragraphs),
-        # "es": ESRetriever(paragraphs),
+    experiments: Dict[str, Experiment] = {
+        "faiss_dpr": Experiment(
+            retriever=FaissRetriever(
+                paragraphs,
+                FaissRetrieverOptions.dpr("./src/models/dpr.faiss")),
+            reader=DprReader()
+        ),
+        "faiss_longformer": Experiment(
+            retriever=FaissRetriever(
+                paragraphs,
+                FaissRetrieverOptions.longformer("./src/models/longformer.faiss")),
+            reader=LongformerReader()
+        ),
+        "es_dpr": Experiment(
+            retriever=ESRetriever(paragraphs),
+            reader=DprReader()
+        ),
+        "es_longformer": Experiment(
+            retriever=ESRetriever(paragraphs),
+            reader=LongformerReader()
+        ),
     }
 
-    for experiment_name, retriever in experiments.items():
-        reader = DprReader()
-
+    for experiment_name, experiment in experiments.items():
+        logger.info(f"Running experiment {experiment_name}...")
         for idx in range(subset_idx):
             question = questions_test["question"][idx]
             answer = questions_test["answer"][idx]
 
-            scores, context = retriever.retrieve(question, 5)
+            retrieve_timer = timeit(f"{experiment_name}.retrieve")
+            t_retrieve = retrieve_timer(experiment.retriever.retrieve)
+
+            read_timer = timeit(f"{experiment_name}.read")
+            t_read = read_timer(experiment.reader.read)
+
+            print(f"\x1b[1K\r[{idx+1:03}] - \"{question}\"", end='')
+
+            scores, context = t_retrieve(question, 5)
             reader_input = context_to_reader_input(context)
 
-            # workaround so we can use the decorator with a dynamic name for time recording
-            time_wrapper = timeit(f"{experiment_name}.read")
-            answers = time_wrapper(reader.read)(question, reader_input, 5)
+            # workaround so we can use the decorator with a dynamic name for
+            # time recording
+            answers = t_read(question, reader_input, 5)
 
             # Calculate softmaxed scores for readable output
-            sm = torch.nn.Softmax(dim=0)
-            document_scores = sm(torch.Tensor(
-                [pred.relevance_score for pred in answers]))
-            span_scores = sm(torch.Tensor(
-                [pred.span_score for pred in answers]))
+            # sm = torch.nn.Softmax(dim=0)
+            # document_scores = sm(torch.Tensor(
+            #     [pred.relevance_score for pred in answers]))
+            # span_scores = sm(torch.Tensor(
+            #     [pred.span_score for pred in answers]))
 
-            print_answers(answers, scores, context)
+            # print_answers(answers, scores, context)
 
             # TODO evaluation and storing of results
+        print()
 
     times = get_times()
-    print(times)
+
+    df = pd.DataFrame(times)
+    os.makedirs("./results/", exist_ok=True)
+    df.to_csv("./results/timings.csv")
+
+
     # TODO evaluation and storing of results
 
     # # Initialize retriever

poetry.lock

@@ -212,6 +212,20 @@ category = "main"
 optional = false
 python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
 
+[[package]]
+name = "coloredlogs"
+version = "15.0.1"
+description = "Colored terminal output for Python's logging module"
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+
+[package.dependencies]
+humanfriendly = ">=9.1"
+
+[package.extras]
+cron = ["capturer (>=2.4)"]
+
 [[package]]
 name = "cryptography"
 version = "36.0.2"
@@ -266,13 +280,13 @@ xxhash = "*"
 apache-beam = ["apache-beam (>=2.26.0)"]
 audio = ["librosa"]
 benchmarks = ["numpy (==1.18.5)", "tensorflow (==2.3.0)", "torch (==1.6.0)", "transformers (==3.0.2)"]
-dev = ["absl-py", "pytest", "pytest-datadir", "pytest-xdist", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "aiobotocore", "boto3", "botocore", "faiss-cpu (>=1.6.4)", "fsspec", "moto[s3,server] (==2.0.4)", "rarfile (>=4.0)", "s3fs (==2021.08.1)", "tensorflow (>=2.3,!=2.6.0,!=2.6.1)", "torch", "torchaudio", "soundfile", "transformers", "bs4", "conllu", "h5py", "langdetect", "lxml", "mwparserfromhell", "nltk", "openpyxl", "py7zr", "tldextract", "zstandard", "bert-score (>=0.3.6)", "rouge-score", "sacrebleu", "scipy", "seqeval", "scikit-learn", "jiwer", "sentencepiece", "torchmetrics (==0.6.0)", "mauve-text", "toml (>=0.10.1)", "requests-file (>=1.5.1)", "tldextract (>=3.1.0)", "texttable (>=1.6.3)", "Werkzeug (>=1.0.1)", "six (>=1.15.0,<1.16.0)", "Pillow (>=6.2.1)", "librosa", "wget (>=3.2)", "pytorch-nlp (==0.5.0)", "pytorch-lightning", "fastBPE (==0.1.0)", "fairseq", "black (>=22.0,<23.0)", "flake8 (>=3.8.3)", "isort (>=5.0.0)", "pyyaml (>=5.3.1)", "importlib-resources"]
+dev = ["absl-py", "pytest", "pytest-datadir", "pytest-xdist", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "aiobotocore", "boto3", "botocore", "faiss-cpu (>=1.6.4)", "fsspec", "moto[server,s3] (==2.0.4)", "rarfile (>=4.0)", "s3fs (==2021.08.1)", "tensorflow (>=2.3,!=2.6.0,!=2.6.1)", "torch", "torchaudio", "soundfile", "transformers", "bs4", "conllu", "h5py", "langdetect", "lxml", "mwparserfromhell", "nltk", "openpyxl", "py7zr", "tldextract", "zstandard", "bert-score (>=0.3.6)", "rouge-score", "sacrebleu", "scipy", "seqeval", "scikit-learn", "jiwer", "sentencepiece", "torchmetrics (==0.6.0)", "mauve-text", "toml (>=0.10.1)", "requests-file (>=1.5.1)", "tldextract (>=3.1.0)", "texttable (>=1.6.3)", "Werkzeug (>=1.0.1)", "six (>=1.15.0,<1.16.0)", "Pillow (>=6.2.1)", "librosa", "wget (>=3.2)", "pytorch-nlp (==0.5.0)", "pytorch-lightning", "fastBPE (==0.1.0)", "fairseq", "black (>=22.0,<23.0)", "flake8 (>=3.8.3)", "isort (>=5.0.0)", "pyyaml (>=5.3.1)", "importlib-resources"]
 docs = ["docutils (==0.16.0)", "recommonmark", "sphinx (==3.1.2)", "sphinx-markdown-tables", "sphinx-rtd-theme (==0.4.3)", "sphinxext-opengraph (==0.4.1)", "sphinx-copybutton", "fsspec (<2021.9.0)", "s3fs", "sphinx-panels", "sphinx-inline-tabs", "myst-parser", "Markdown (!=3.3.5)"]
 quality = ["black (>=22.0,<23.0)", "flake8 (>=3.8.3)", "isort (>=5.0.0)", "pyyaml (>=5.3.1)"]
 s3 = ["fsspec", "boto3", "botocore", "s3fs"]
 tensorflow = ["tensorflow (>=2.2.0,!=2.6.0,!=2.6.1)"]
 tensorflow_gpu = ["tensorflow-gpu (>=2.2.0,!=2.6.0,!=2.6.1)"]
-tests = ["absl-py", "pytest", "pytest-datadir", "pytest-xdist", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "aiobotocore", "boto3", "botocore", "faiss-cpu (>=1.6.4)", "fsspec", "moto[s3,server] (==2.0.4)", "rarfile (>=4.0)", "s3fs (==2021.08.1)", "tensorflow (>=2.3,!=2.6.0,!=2.6.1)", "torch", "torchaudio", "soundfile", "transformers", "bs4", "conllu", "h5py", "langdetect", "lxml", "mwparserfromhell", "nltk", "openpyxl", "py7zr", "tldextract", "zstandard", "bert-score (>=0.3.6)", "rouge-score", "sacrebleu", "scipy", "seqeval", "scikit-learn", "jiwer", "sentencepiece", "torchmetrics (==0.6.0)", "mauve-text", "toml (>=0.10.1)", "requests-file (>=1.5.1)", "tldextract (>=3.1.0)", "texttable (>=1.6.3)", "Werkzeug (>=1.0.1)", "six (>=1.15.0,<1.16.0)", "Pillow (>=6.2.1)", "librosa", "wget (>=3.2)", "pytorch-nlp (==0.5.0)", "pytorch-lightning", "fastBPE (==0.1.0)", "fairseq", "importlib-resources"]
+tests = ["absl-py", "pytest", "pytest-datadir", "pytest-xdist", "apache-beam (>=2.26.0)", "elasticsearch (<8.0.0)", "aiobotocore", "boto3", "botocore", "faiss-cpu (>=1.6.4)", "fsspec", "moto[server,s3] (==2.0.4)", "rarfile (>=4.0)", "s3fs (==2021.08.1)", "tensorflow (>=2.3,!=2.6.0,!=2.6.1)", "torch", "torchaudio", "soundfile", "transformers", "bs4", "conllu", "h5py", "langdetect", "lxml", "mwparserfromhell", "nltk", "openpyxl", "py7zr", "tldextract", "zstandard", "bert-score (>=0.3.6)", "rouge-score", "sacrebleu", "scipy", "seqeval", "scikit-learn", "jiwer", "sentencepiece", "torchmetrics (==0.6.0)", "mauve-text", "toml (>=0.10.1)", "requests-file (>=1.5.1)", "tldextract (>=3.1.0)", "texttable (>=1.6.3)", "Werkzeug (>=1.0.1)", "six (>=1.15.0,<1.16.0)", "Pillow (>=6.2.1)", "librosa", "wget (>=3.2)", "pytorch-nlp (==0.5.0)", "pytorch-lightning", "fastBPE (==0.1.0)", "fairseq", "importlib-resources"]
 torch = ["torch"]
 vision = ["Pillow (>=6.2.1)"]
 
@@ -439,7 +453,7 @@ python-versions = ">=3.7"
 
 [[package]]
 name = "fsspec"
-version = "2022.2.0"
+version = "2022.3.0"
 description = "File-system specification"
 category = "main"
 optional = false
@@ -470,10 +484,11 @@ s3 = ["s3fs"]
 sftp = ["paramiko"]
 smb = ["smbprotocol"]
 ssh = ["paramiko"]
+tqdm = ["tqdm"]
 
 [[package]]
 name = "gradio"
-version = "2.9.0"
+version = "2.9.1"
 description = "Python library for easily interacting with trained machine learning models"
 category = "main"
 optional = false
@@ -529,6 +544,17 @@ all = ["pytest", "datasets", "black (>=20.8b1)", "isort (>=5.5.4)", "flake8 (>=3
 dev = ["pytest", "datasets", "black (>=20.8b1)", "isort (>=5.5.4)", "flake8 (>=3.8.3)"]
 quality = ["black (>=20.8b1)", "isort (>=5.5.4)", "flake8 (>=3.8.3)"]
 
+[[package]]
+name = "humanfriendly"
+version = "10.0"
+description = "Human friendly output for text interfaces using Python"
+category = "main"
+optional = false
+python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*"
+
+[package.dependencies]
+pyreadline3 = {version = "*", markers = "sys_platform == \"win32\" and python_version >= \"3.8\""}
+
 [[package]]
 name = "idna"
 version = "3.3"
@@ -1099,6 +1125,14 @@ python-versions = ">=3.6"
 [package.extras]
 diagrams = ["jinja2", "railroad-diagrams"]
 
+[[package]]
+name = "pyreadline3"
+version = "3.4.1"
+description = "A python implementation of GNU readline."
+category = "main"
+optional = false
+python-versions = "*"
+
 [[package]]
 name = "python-dateutil"
 version = "2.8.2"
@@ -1498,7 +1532,7 @@ multidict = ">=4.0"
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.8"
-content-hash = "a9ce48f30c8568321f3f4576e1c4987ef94a4216201ba4bce2dc719c397d5da6"
+content-hash = "881ba67f914b3c0690bcb34810061252ee77cebc0dac49b5ae76348394d810a8"
 
 [metadata.files]
 aiohttp = [
@@ -1699,6 +1733,10 @@ colorama = [
     {file = "colorama-0.4.4-py2.py3-none-any.whl", hash = "sha256:9f47eda37229f68eee03b24b9748937c7dc3868f906e8ba69fbcbdd3bc5dc3e2"},
     {file = "colorama-0.4.4.tar.gz", hash = "sha256:5941b2b48a20143d2267e95b1c2a7603ce057ee39fd88e7329b0c292aa16869b"},
 ]
+coloredlogs = [
+    {file = "coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934"},
+    {file = "coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0"},
+]
 cryptography = [
     {file = "cryptography-36.0.2-cp36-abi3-macosx_10_10_universal2.whl", hash = "sha256:4e2dddd38a5ba733be6a025a1475a9f45e4e41139d1321f412c6b360b19070b6"},
     {file = "cryptography-36.0.2-cp36-abi3-macosx_10_10_x86_64.whl", hash = "sha256:4881d09298cd0b669bb15b9cfe6166f16fc1277b4ed0d04a22f3d6430cb30f1d"},
@@ -1880,12 +1918,12 @@ frozenlist = [
     {file = "frozenlist-1.3.0.tar.gz", hash = "sha256:ce6f2ba0edb7b0c1d8976565298ad2deba6f8064d2bebb6ffce2ca896eb35b0b"},
 ]
 fsspec = [
-    {file = "fsspec-2022.2.0-py3-none-any.whl", hash = "sha256:eb9c9d9aee49d23028deefffe53e87c55d3515512c63f57e893710301001449a"},
-    {file = "fsspec-2022.2.0.tar.gz", hash = "sha256:20322c659538501f52f6caa73b08b2ff570b7e8ea30a86559721d090e473ad5c"},
+    {file = "fsspec-2022.3.0-py3-none-any.whl", hash = "sha256:a53491b003210fce6911dd8f2d37e20c41a27ce52a655eef11b885d1578ed4cf"},
+    {file = "fsspec-2022.3.0.tar.gz", hash = "sha256:fd582cc4aa0db5968bad9317cae513450eddd08b2193c4428d9349265a995523"},
 ]
 gradio = [
-    {file = "gradio-2.9.0-py3-none-any.whl", hash = "sha256:02c3604d8c662dc35a60e75f55c3de175f8e2c30bf868c39e82f8c20a608d80b"},
-    {file = "gradio-2.9.0.tar.gz", hash = "sha256:2cfbde23425c97959291d88ceae55e3d83e1a32915a0e9f7032c8c81bd4f5b63"},
+    {file = "gradio-2.9.1-py3-none-any.whl", hash = "sha256:877616dcda82e0e13bc04404c13f084c7b3a06cccc314a4db06b21c5f15f6190"},
+    {file = "gradio-2.9.1.tar.gz", hash = "sha256:d9dfde81f064f38bcd95967316501ab40698fec0bcc4435dd00ea4578f695042"},
 ]
 h11 = [
     {file = "h11-0.13.0-py3-none-any.whl", hash = "sha256:8ddd78563b633ca55346c8cd41ec0af27d3c79931828beffb46ce70a379e7442"},
@@ -1895,6 +1933,10 @@ huggingface-hub = [
     {file = "huggingface_hub-0.4.0-py3-none-any.whl", hash = "sha256:808021af1ce1111104973ae54d81738eaf40be6d1e82fc6bdedb82f81c6206e7"},
     {file = "huggingface_hub-0.4.0.tar.gz", hash = "sha256:f0e3389f8988eb7781b17de520ae7fd0aa50d9823534e3ae55344d943a88ac87"},
 ]
+humanfriendly = [
+    {file = "humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477"},
+    {file = "humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc"},
+]
 idna = [
     {file = "idna-3.3-py3-none-any.whl", hash = "sha256:84d9dd047ffa80596e0f246e2eab0b391788b0503584e8945f2368256d2735ff"},
     {file = "idna-3.3.tar.gz", hash = "sha256:9d643ff0a55b762d5cdb124b8eaa99c66322e2157b69160bc32796e824360e6d"},
@@ -2193,6 +2235,7 @@ numpy = [
     {file = "numpy-1.22.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8251ed96f38b47b4295b1ae51631de7ffa8260b5b087808ef09a39a9d66c97ab"},
     {file = "numpy-1.22.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:48a3aecd3b997bf452a2dedb11f4e79bc5bfd21a1d4cc760e703c31d57c84b3e"},
     {file = "numpy-1.22.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a3bae1a2ed00e90b3ba5f7bd0a7c7999b55d609e0c54ceb2b076a25e345fa9f4"},
+    {file = "numpy-1.22.3-cp310-cp310-win32.whl", hash = "sha256:f950f8845b480cffe522913d35567e29dd381b0dc7e4ce6a4a9f9156417d2430"},
     {file = "numpy-1.22.3-cp310-cp310-win_amd64.whl", hash = "sha256:08d9b008d0156c70dc392bb3ab3abb6e7a711383c3247b410b39962263576cd4"},
     {file = "numpy-1.22.3-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:201b4d0552831f7250a08d3b38de0d989d6f6e4658b709a02a73c524ccc6ffce"},
     {file = "numpy-1.22.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:f8c1f39caad2c896bc0018f699882b345b2a63708008be29b1f355ebf6f933fe"},
@@ -2510,6 +2553,10 @@ pyparsing = [
     {file = "pyparsing-3.0.7-py3-none-any.whl", hash = "sha256:a6c06a88f252e6c322f65faf8f418b16213b51bdfaece0524c1c1bc30c63c484"},
     {file = "pyparsing-3.0.7.tar.gz", hash = "sha256:18ee9022775d270c55187733956460083db60b37d0d0fb357445f3094eed3eea"},
 ]
+pyreadline3 = [
+    {file = "pyreadline3-3.4.1-py3-none-any.whl", hash = "sha256:b0efb6516fd4fb07b45949053826a62fa4cb353db5be2bbb4a7aa1fdd1e345fb"},
+    {file = "pyreadline3-3.4.1.tar.gz", hash = "sha256:6f3d1f7b8a31ba32b73917cefc1f28cc660562f39aea8646d30bd6eff21f7bae"},
+]
 python-dateutil = [
     {file = "python-dateutil-2.8.2.tar.gz", hash = "sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86"},
     {file = "python_dateutil-2.8.2-py2.py3-none-any.whl", hash = "sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9"},

pyproject.toml

@@ -15,6 +15,7 @@ python-dotenv = "^0.19.2"
 elasticsearch = "^8.1.0"
 gradio = {extras = ["Jinja2"], version = "^2.9.0"}
 Jinja2 = "^3.1.1"
+coloredlogs = "^15.0.1"
 
 [tool.poetry.dev-dependencies]
 flake8 = "^4.0.1"

query.py

@@ -16,7 +16,12 @@ from src.retrievers.faiss_retriever import (
     FaissRetrieverOptions
 )
 from src.utils.preprocessing import context_to_reader_input
-from src.utils.log import get_logger
+from src.utils.log import logger
+
+
+# Setup environment
+load_dotenv()
+transformers.logging.set_verbosity_error()
 
 
 def get_retriever(paragraphs: DatasetDict,
@@ -123,11 +128,6 @@ def main(args: argparse.Namespace):
 
 
 if __name__ == "__main__":
-    # Setup environment
-    load_dotenv()
-    logger = get_logger()
-    transformers.logging.set_verbosity_error()
-
     # Set up CLI arguments
     parser = argparse.ArgumentParser(
         formatter_class=argparse.MetavarTypeHelpFormatter

results/timings.csv

@@ -0,0 +1,60 @@
+,faiss_dpr.retrieve,faiss_dpr.read,faiss_longformer.retrieve,faiss_longformer.read,es_dpr.retrieve,es_dpr.read,es_longformer.retrieve,es_longformer.read
+0,0.30384302139282227,4.566400051116943,0.9227948188781738,5.768368244171143,0.01930093765258789,2.7453649044036865,0.010576009750366211,4.998417854309082
+1,0.04573678970336914,1.9288370609283447,0.8380529880523682,5.916611671447754,0.018373966217041016,1.4845240116119385,0.012102842330932617,5.1692070960998535
+2,0.04764819145202637,0.6628780364990234,0.7756149768829346,5.4998250007629395,0.015324831008911133,1.7706871032714844,0.012642860412597656,5.202448844909668
+3,0.04507589340209961,1.219634771347046,0.8142738342285156,5.726102113723755,0.021118879318237305,1.987663984298706,0.012515068054199219,5.1083409786224365
+4,0.04347515106201172,1.5222840309143066,0.814906120300293,5.672412872314453,0.013732194900512695,1.660247802734375,0.011805057525634766,5.313212156295776
+5,0.07470989227294922,1.5599188804626465,0.8422539234161377,5.75390100479126,0.018023014068603516,1.5782928466796875,0.013046741485595703,5.419210195541382
+6,0.06162095069885254,1.4178202152252197,0.7837569713592529,4.765166282653809,0.014074325561523438,0.7626080513000488,0.010712146759033203,4.976129055023193
+7,0.0451970100402832,1.134779691696167,0.7723889350891113,4.5784592628479,0.012156963348388672,1.959972858428955,0.011015892028808594,4.5342161655426025
+8,0.03589582443237305,0.8912148475646973,0.8142461776733398,5.212156295776367,0.009800195693969727,1.820624828338623,0.009167194366455078,4.468229293823242
+9,0.06033587455749512,0.37888431549072266,1.1510162353515625,5.395290851593018,0.015815258026123047,1.0247371196746826,0.010970830917358398,5.1864588260650635
+10,0.056854963302612305,0.6068317890167236,0.7839999198913574,4.668170928955078,0.013895988464355469,0.8482949733734131,0.011837005615234375,4.493913173675537
+11,0.04697012901306152,0.341174840927124,0.8089311122894287,5.535298109054565,0.01624298095703125,1.4673452377319336,0.01172780990600586,5.264245986938477
+12,0.0444178581237793,1.4774150848388672,0.7612121105194092,5.504917860031128,0.01690196990966797,2.42773699760437,0.011591196060180664,4.463428974151611
+13,0.06343889236450195,0.8000409603118896,0.9072589874267578,6.015661954879761,0.010895252227783203,2.21577787399292,0.012058019638061523,5.302258253097534
+14,0.05022692680358887,0.9474368095397949,0.8324599266052246,6.4684131145477295,0.016222000122070312,3.1302390098571777,0.010799169540405273,5.240310907363892
+15,0.08313488960266113,0.746314287185669,0.8373219966888428,6.741006851196289,0.017467975616455078,1.353593111038208,0.020203113555908203,5.192620038986206
+16,0.17216706275939941,1.136448860168457,0.7760000228881836,5.48329496383667,0.018398046493530273,0.5325403213500977,0.012536287307739258,5.264941215515137
+17,0.04575324058532715,0.5853927135467529,0.7855441570281982,5.960904121398926,0.019134998321533203,2.6092309951782227,0.012146949768066406,5.331835031509399
+18,0.04746294021606445,0.9219169616699219,0.9516820907592773,10.146074295043945,0.011114835739135742,2.220487117767334,0.011821985244750977,5.161020994186401
+19,0.0443730354309082,0.4667840003967285,1.3496372699737549,8.213943719863892,0.01042485237121582,2.841907024383545,0.011536121368408203,5.236280918121338
+20,0.06004190444946289,0.6129250526428223,1.3677341938018799,7.00742769241333,0.022186994552612305,1.6846930980682373,0.010824918746948242,5.377984046936035
+21,0.06920814514160156,0.6232960224151611,1.4656860828399658,6.424375057220459,0.011613845825195312,1.0811800956726074,0.014858007431030273,5.279160022735596
+22,0.04999184608459473,0.6539132595062256,0.8720510005950928,5.889069080352783,0.016654014587402344,1.6599159240722656,0.012172698974609375,5.177525043487549
+23,0.05750322341918945,1.0169367790222168,0.9728169441223145,6.934185028076172,0.01772904396057129,1.2837882041931152,0.011108160018920898,5.186945199966431
+24,0.06264281272888184,1.7151312828063965,1.3927390575408936,7.122100114822388,0.016143798828125,1.5387201309204102,0.011415958404541016,4.558846950531006
+25,0.04831504821777344,0.7839398384094238,1.1007087230682373,5.4652369022369385,0.01099395751953125,1.5678913593292236,0.011976242065429688,4.612828969955444
+26,0.048091888427734375,0.9228200912475586,0.8567941188812256,4.832158803939819,0.013817787170410156,2.0290918350219727,0.015846967697143555,4.845104217529297
+27,0.04568672180175781,0.8964569568634033,0.7873432636260986,4.592561960220337,0.010241985321044922,0.3145887851715088,0.014873743057250977,4.759660720825195
+28,0.04340720176696777,0.5004391670227051,0.8122010231018066,4.68702507019043,0.012717008590698242,0.9207170009613037,0.014780759811401367,4.955734968185425
+29,0.045496225357055664,2.106112003326416,0.7901277542114258,5.48145604133606,0.009778976440429688,1.1795310974121094,0.011364936828613281,5.397800922393799
+30,0.05589914321899414,3.3801350593566895,0.7913417816162109,4.76953387260437,0.01169896125793457,2.8297739028930664,0.012899160385131836,4.7149817943573
+31,0.038469791412353516,1.2037632465362549,0.812114953994751,4.819751977920532,0.010591983795166016,1.0633080005645752,0.011631011962890625,4.603592157363892
+32,0.043640851974487305,0.7455379962921143,0.7684001922607422,5.490149021148682,0.010446786880493164,1.509342908859253,0.01111912727355957,5.431332111358643
+33,0.0411829948425293,0.7775781154632568,0.7725949287414551,5.5284202098846436,0.011181116104125977,1.4173851013183594,0.01881098747253418,5.2474939823150635
+34,0.04268312454223633,1.3576858043670654,0.7971670627593994,5.488955974578857,0.016661882400512695,0.6669139862060547,0.011193990707397461,5.231971263885498
+35,0.0432438850402832,0.49681520462036133,0.7736399173736572,4.675936698913574,0.013994932174682617,0.7481560707092285,0.01053619384765625,4.871787071228027
+36,0.038790225982666016,1.9925789833068848,0.7900221347808838,4.716547012329102,0.010754108428955078,0.8104310035705566,0.011471986770629883,4.582187175750732
+37,0.04674410820007324,0.8766942024230957,0.8192441463470459,5.454381704330444,0.012632131576538086,3.3098862171173096,0.01573491096496582,5.5617289543151855
+38,0.04983806610107422,0.5784440040588379,0.768744945526123,5.399757146835327,0.017091035842895508,1.0388100147247314,0.020289897918701172,5.327627897262573
+39,0.039936065673828125,0.9906370639801025,0.7951750755310059,4.816935062408447,0.009315729141235352,0.8949270248413086,0.012948989868164062,4.823601007461548
+40,0.04812121391296387,5.3651018142700195,0.7833847999572754,4.673122882843018,0.010359048843383789,1.6986067295074463,0.012405872344970703,4.822720050811768
+41,0.037177085876464844,0.8579537868499756,0.768902063369751,4.705405950546265,0.01087808609008789,1.1154420375823975,0.009827136993408203,5.295310020446777
+42,0.03615593910217285,0.6045210361480713,0.7767770290374756,4.721595048904419,0.012170076370239258,1.168515920639038,0.014606952667236328,4.778914213180542
+43,0.04032111167907715,1.0840678215026855,0.8039369583129883,5.5514678955078125,0.011640071868896484,3.7264089584350586,0.015080928802490234,6.431236028671265
+44,0.12291288375854492,2.7860946655273438,0.7999370098114014,4.700652122497559,0.010669708251953125,3.5256330966949463,0.00997614860534668,5.010454893112183
+45,0.03981208801269531,0.8575420379638672,0.7781379222869873,4.649600028991699,0.011057853698730469,3.4576022624969482,0.011123895645141602,5.414888143539429
+46,0.046558380126953125,0.6096041202545166,0.839914083480835,5.3846352100372314,0.0264890193939209,3.282578945159912,0.013241052627563477,6.356001853942871
+47,0.044730186462402344,0.6428439617156982,0.7774860858917236,5.471776962280273,0.009460926055908203,3.3428800106048584,0.012679100036621094,5.476663112640381
+48,0.04798007011413574,1.3710291385650635,0.7838289737701416,5.5646140575408936,0.011425018310546875,1.5621020793914795,0.019647836685180664,5.403181076049805
+49,0.06305599212646484,1.7375829219818115,0.7764248847961426,5.582126140594482,0.010413169860839844,1.6502351760864258,0.011098146438598633,6.15350604057312
+50,0.04781007766723633,0.919248104095459,0.8292880058288574,4.79367995262146,0.01233983039855957,4.761476755142212,0.01306009292602539,4.901428937911987
+51,0.04294776916503906,0.9060940742492676,0.7503399848937988,4.69527006149292,0.010550737380981445,1.3250057697296143,0.012276887893676758,4.790279388427734
+52,0.0449681282043457,0.74688720703125,0.7592051029205322,4.672075986862183,0.010587215423583984,1.7173192501068115,0.012059926986694336,4.9025468826293945
+53,0.04381895065307617,1.2078793048858643,0.8653321266174316,4.4878456592559814,0.008989810943603516,4.782422065734863,0.012001752853393555,5.331949949264526
+54,0.06584286689758301,1.0724549293518066,0.7348787784576416,5.094892740249634,0.00992584228515625,1.6900959014892578,0.018785953521728516,6.253833293914795
+55,0.05147194862365723,2.172264337539673,0.7367160320281982,5.056357145309448,0.009489059448242188,2.6061501502990723,0.011726140975952148,6.203222036361694
+56,0.05095195770263672,2.0959391593933105,0.7292170524597168,5.11798882484436,0.011085987091064453,1.258976936340332,0.02020883560180664,6.149781942367554
+57,0.05691885948181152,0.7286462783813477,0.7636628150939941,5.1169517040252686,0.01094675064086914,0.7379579544067383,0.012721061706542969,5.3568830490112305
+58,0.04192709922790527,0.8154990673065186,0.7308712005615234,5.066887140274048,0.010490894317626953,2.724623203277588,0.01871800422668457,5.368160009384155

src/reader.py

@@ -1,2 +0,0 @@
-class Reader():
-    pass

src/readers/dpr_reader.py

@@ -1,9 +1,13 @@
 from transformers import DPRReader, DPRReaderTokenizer
 from typing import List, Dict, Tuple
+from dotenv import load_dotenv
 
 from src.readers.base_reader import Reader
 
 
+load_dotenv()
+
+
 class DprReader(Reader):
     def __init__(self) -> None:
         self._tokenizer = DPRReaderTokenizer.from_pretrained(

src/readers/longformer_reader.py

@@ -1,17 +1,21 @@
 import torch
 from transformers import (
-    LongformerTokenizerFast,
+    LongformerTokenizer,
     LongformerForQuestionAnswering
 )
 from typing import List, Dict, Tuple
+from dotenv import load_dotenv
 
 from src.readers.base_reader import Reader
 
 
+load_dotenv()
+
+
 class LongformerReader(Reader):
     def __init__(self) -> None:
         checkpoint = "valhalla/longformer-base-4096-finetuned-squadv1"
-        self.tokenizer = LongformerTokenizerFast.from_pretrained(checkpoint)
+        self.tokenizer = LongformerTokenizer.from_pretrained(checkpoint)
         self.model = LongformerForQuestionAnswering.from_pretrained(checkpoint)
 
     def read(self,
@@ -21,8 +25,7 @@ class LongformerReader(Reader):
         answers = []
 
         for text in context['texts']:
-            encoding = self.tokenizer(
-                query, text, return_tensors="pt")
+            encoding = self.tokenizer(query, text, return_tensors="pt")
             input_ids = encoding["input_ids"]
             attention_mask = encoding["attention_mask"]
             outputs = self.model(input_ids, attention_mask=attention_mask)

src/retrievers/es_retriever.py

@@ -1,13 +1,17 @@
+import imp
 import os
 
 from datasets import DatasetDict
 from elasticsearch import Elasticsearch
+from elastic_transport import ConnectionError
+from dotenv import load_dotenv
 
 from src.retrievers.base_retriever import RetrieveType, Retriever
-from src.utils.log import get_logger
+from src.utils.log import logger
 from src.utils.timing import timeit
 
-logger = get_logger()
+
+load_dotenv()
 
 
 class ESRetriever(Retriever):
@@ -23,6 +27,13 @@ class ESRetriever(Retriever):
             http_auth=(es_username, es_password),
             ca_certs="./http_ca.crt")
 
+        try:
+            self.client.info()
+        except ConnectionError:
+            logger.error("Could not connect to ElasticSearch. " + 
+                         "Make sure it is running. Exiting now...")
+            exit()
+
         if self.client.indices.exists(index="paragraphs"):
             self.paragraphs.load_elasticsearch_index(
                 "paragraphs", es_index_name="paragraphs",
@@ -34,6 +45,5 @@ class ESRetriever(Retriever):
                                                     es_index_name="paragraphs",
                                                     es_client=self.client)
 
-    @timeit("esretriever.retrieve")
     def retrieve(self, query: str, k: int = 5) -> RetrieveType:
         return self.paragraphs.get_nearest_examples("paragraphs", query, k)

src/retrievers/faiss_retriever.py

@@ -2,6 +2,7 @@ import os
 import os.path
 import torch
 
+from dotenv import load_dotenv
 from datasets import DatasetDict
 from dataclasses import dataclass
 from transformers import (
@@ -10,22 +11,18 @@ from transformers import (
     DPRQuestionEncoder,
     DPRQuestionEncoderTokenizerFast,
     LongformerModel,
-    LongformerTokenizerFast
+    LongformerTokenizer
 )
 from transformers.modeling_utils import PreTrainedModel
 from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
 
 from src.retrievers.base_retriever import RetrieveType, Retriever
-from src.utils.log import get_logger
+from src.utils.log import logger
 from src.utils.preprocessing import remove_formulas
 from src.utils.timing import timeit
 
-# Hacky fix for FAISS error on macOS
-# See https://stackoverflow.com/a/63374568/4545692
-os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
 
-
-logger = get_logger()
+load_dotenv()
 
 
 @dataclass
@@ -59,10 +56,10 @@ class FaissRetrieverOptions:
     @staticmethod
     def longformer(embedding_path: str):
         encoder = LongformerModel.from_pretrained(
-            "allenai/longformer-base-4096"
+            "valhalla/longformer-base-4096-finetuned-squadv1"
         )
-        tokenizer = LongformerTokenizerFast.from_pretrained(
-            "allenai/longformer-base-4096"
+        tokenizer = LongformerTokenizer.from_pretrained(
+            "valhalla/longformer-base-4096-finetuned-squadv1"
         )
         return FaissRetrieverOptions(
             ctx_encoder=encoder,
@@ -145,7 +142,6 @@ class FaissRetriever(Retriever):
 
             return index
 
-    @timeit("faissretriever.retrieve")
     def retrieve(self, query: str, k: int = 5) -> RetrieveType:
         question_embedding = self._embed_question(query)
         scores, results = self.index.get_nearest_examples(

src/utils/log.py

@@ -1,3 +1,4 @@
+import coloredlogs
 import logging
 import os
 
@@ -5,27 +6,27 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
+# creates a default logger for the project. We declare it in the global scope
+# so it acts like a singleton
+logger = logging.getLogger("Flashcards")
 
-def get_logger():
-    # creates a default logger for the project
-    logger = logging.getLogger("Flashcards")
+log_level = os.getenv("LOG_LEVEL", "INFO")
+logger.setLevel(log_level)
 
-    log_level = os.getenv("LOG_LEVEL", "INFO")
-    logger.setLevel(log_level)
+# Log format
+formatter = coloredlogs.ColoredFormatter(
+    "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
 
-    # Log format
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+# stout
+ch = logging.StreamHandler()
+ch.setFormatter(formatter)
 
-    # file handler
-    fh = logging.FileHandler("logs.log")
-    fh.setFormatter(formatter)
+# colored output so log messages stand out more
+# coloredlogs.install(level=log_level, logger=logger)
 
-    # stout
-    ch = logging.StreamHandler()
-    ch.setFormatter(formatter)
+# file handler
+fh = logging.FileHandler("logs.log")
+fh.setFormatter(formatter)
 
-    logger.addHandler(fh)
-    logger.addHandler(ch)
-
-    return logger
+logger.addHandler(fh)
+logger.addHandler(ch)

src/utils/timing.py

@@ -1,11 +1,8 @@
 import time
-from typing import Dict
+from typing import Dict, List
 from dotenv import load_dotenv
 import os
-from src.utils.log import get_logger
-
-
-logger = get_logger()
+from src.utils.log import logger
 
 
 load_dotenv()
@@ -17,7 +14,7 @@ if ENABLE_TIMING:
     logger.info("Timing is enabled")
 
 
-TimingType = Dict[str, float]
+TimingType = Dict[str, List[float]]
 
 TIMES: TimingType = {}