add experiment code

README.old.md

@@ -6,12 +6,12 @@
   - [ ] Formules enzo eruit filteren
   - [ ] Splitsen op zinnen...?
 - [ ] Meer language models proberen
-- [ ] Elasticsearch
-- [ ] CLI voor vragen beantwoorden
+- [X] Elasticsearch
+- [X] CLI voor vragen beantwoorden
 
 ### Extra dingen
 
-- [ ] Huggingface spaces demo
+- [X] Huggingface spaces demo
 - [ ] Question generation voor finetuning
 - [ ] Language model finetunen
 

main.py

@@ -1,19 +1,20 @@
-import os
 import random
-from typing import cast
-import time
+from typing import Dict, cast
 
 import torch
 import transformers
 from datasets import DatasetDict, load_dataset
 from dotenv import load_dotenv
+from query import print_answers
 
 from src.evaluation import evaluate
 from src.readers.dpr_reader import DprReader
+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.utils.preprocessing import context_to_reader_input
+from src.utils.timing import get_times, timeit
 
 logger = get_logger()
 
@@ -26,62 +27,97 @@ if __name__ == '__main__':
         "GroNLP/ik-nlp-22_slp", "paragraphs"))
     questions = cast(DatasetDict, load_dataset(dataset_name, "questions"))
 
-    questions_test = questions["test"]
-
-    # Initialize retriever
-    retriever = FaissRetriever(paragraphs)
-    #retriever = ESRetriever(paragraphs)
-
-    # Retrieve example
-    # random.seed(111)
-    random_index = random.randint(0, len(questions_test["question"])-1)
-    example_q = questions_test["question"][random_index]
-    example_a = questions_test["answer"][random_index]
-
-    scores, result = retriever.retrieve(example_q)
-    reader_input = context_to_reader_input(result)
-
-    # TODO: use new code from query.py to clean this up
-    # Initialize reader
-    reader = DprReader()
-    answers = reader.read(example_q, reader_input)
-
-    # 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]))
-
-    print(example_q)
-    for answer_i, answer in enumerate(answers):
-        print(f"[{answer_i + 1}]: {answer.text}")
-        print(f"\tDocument {answer.doc_id}", end='')
-        print(f"\t(score {document_scores[answer_i] * 100:.02f})")
-        print(f"\tSpan {answer.start_index}-{answer.end_index}", end='')
-        print(f"\t(score {span_scores[answer_i] * 100:.02f})")
-        print()  # Newline
-
-    # print(f"Example q: {example_q} answer: {result['text'][0]}")
-
-    # for i, score in enumerate(scores):
-    #     print(f"Result {i+1} (score: {score:.02f}):")
-    #     print(result['text'][i])
-
-    # Determine best answer we want to evaluate
-    highest, highest_index = 0, 0
-    for i, value in enumerate(span_scores):
-        if value + document_scores[i] > highest:
-            highest = value + document_scores[i]
-            highest_index = i
-
-    # Retrieve exact match and F1-score
-    exact_match, f1_score = evaluate(
-        example_a, answers[highest_index].text)
-    print(f"Gold answer: {example_a}\n"
-          f"Predicted answer: {answers[highest_index].text}\n"
-          f"Exact match: {exact_match:.02f}\n"
-          f"F1-score: {f1_score:.02f}")
+    # Only doing a few questions for speed
+    subset_idx = 3
+    questions_test = questions["test"][:subset_idx]
+
+    experiments: Dict[str, Retriever] = {
+        "faiss": FaissRetriever(paragraphs),
+        # "es": ESRetriever(paragraphs),
+    }
+
+    for experiment_name, retriever in experiments.items():
+        reader = DprReader()
+
+        for idx in range(subset_idx):
+            question = questions_test["question"][idx]
+            answer = questions_test["answer"][idx]
+
+            scores, context = retriever.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)
+
+            # 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]))
+
+            print_answers(answers, scores, context)
+
+            # TODO evaluation and storing of results
+
+    times = get_times()
+    print(times)
+    # TODO evaluation and storing of results
+
+    # # Initialize retriever
+    # retriever = FaissRetriever(paragraphs)
+    # # retriever = ESRetriever(paragraphs)
+
+    # # Retrieve example
+    # # random.seed(111)
+    # random_index = random.randint(0, len(questions_test["question"])-1)
+    # example_q = questions_test["question"][random_index]
+    # example_a = questions_test["answer"][random_index]
+
+    # scores, result = retriever.retrieve(example_q)
+    # reader_input = context_to_reader_input(result)
+
+    # # TODO: use new code from query.py to clean this up
+    # # Initialize reader
+    # answers = reader.read(example_q, reader_input)
+
+    # # 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]))
+
+    # print(example_q)
+    # for answer_i, answer in enumerate(answers):
+    #     print(f"[{answer_i + 1}]: {answer.text}")
+    #     print(f"\tDocument {answer.doc_id}", end='')
+    #     print(f"\t(score {document_scores[answer_i] * 100:.02f})")
+    #     print(f"\tSpan {answer.start_index}-{answer.end_index}", end='')
+    #     print(f"\t(score {span_scores[answer_i] * 100:.02f})")
+    #     print()  # Newline
+
+    # # print(f"Example q: {example_q} answer: {result['text'][0]}")
+
+    # # for i, score in enumerate(scores):
+    # #     print(f"Result {i+1} (score: {score:.02f}):")
+    # #     print(result['text'][i])
+
+    # # Determine best answer we want to evaluate
+    # highest, highest_index = 0, 0
+    # for i, value in enumerate(span_scores):
+    #     if value + document_scores[i] > highest:
+    #         highest = value + document_scores[i]
+    #         highest_index = i
+
+    # # Retrieve exact match and F1-score
+    # exact_match, f1_score = evaluate(
+    #     example_a, answers[highest_index].text)
+    # print(f"Gold answer: {example_a}\n"
+    #       f"Predicted answer: {answers[highest_index].text}\n"
+    #       f"Exact match: {exact_match:.02f}\n"
+    #       f"F1-score: {f1_score:.02f}")
 
     # Calculate overall performance
     # total_f1 = 0

src/retrievers/base_retriever.py

@@ -1,3 +1,12 @@
+from typing import Dict, List, Tuple
+
+import numpy as np
+
+RetrieveTypeResult = Dict[str, List[str]]
+RetrieveTypeScores = np.ndarray
+RetrieveType = Tuple[RetrieveTypeScores, RetrieveTypeResult]
+
+
 class Retriever():
-    def retrieve(self, query: str, k: int):
-        pass
+    def retrieve(self, query: str, k: int) -> RetrieveType:
+        raise NotImplementedError()

src/retrievers/es_retriever.py

@@ -1,8 +1,11 @@
+import os
+
 from datasets import DatasetDict
-from src.utils.log import get_logger
-from src.retrievers.base_retriever import Retriever
 from elasticsearch import Elasticsearch
-import os
+
+from src.retrievers.base_retriever import RetrieveType, Retriever
+from src.utils.log import get_logger
+from src.utils.timing import timeit
 
 logger = get_logger()
 
@@ -31,5 +34,6 @@ class ESRetriever(Retriever):
                                                     es_index_name="paragraphs",
                                                     es_client=self.client)
 
-    def retrieve(self, query: str, k: int = 5):
+    @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

@@ -10,9 +10,10 @@ from transformers import (
     DPRQuestionEncoderTokenizer,
 )
 
-from src.retrievers.base_retriever import Retriever
+from src.retrievers.base_retriever import RetrieveType, Retriever
 from src.utils.log import get_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
@@ -83,7 +84,8 @@ class FaissRetriever(Retriever):
 
             return index
 
-    def retrieve(self, query: str, k: int = 50):
+    @timeit("faissretriever.retrieve")
+    def retrieve(self, query: str, k: int = 5) -> RetrieveType:
         def embed(q):
             # Inline helper function to perform embedding
             tok = self.q_tokenizer(q, return_tensors="pt", truncation=True)

test.py

@@ -0,0 +1,20 @@
+# %%
+from datasets import load_dataset
+from src.retrievers.faiss_retriever import FaissRetriever
+
+
+data = load_dataset("GroNLP/ik-nlp-22_slp", "paragraphs")
+
+# # %%
+# x = data["test"][:3]
+
+# # %%
+# for y in x:
+
+#     print(y)
+# # %%
+# x.num_rows
+
+# # %%
+retriever = FaissRetriever(data)
+scores, result = retriever.retrieve("hello world")