refactor evaluation

main.py

@@ -1,23 +1,38 @@
-from src.fais_retriever import FAISRetriever
-from src.utils.log import get_logger
+from datasets import DatasetDict, load_dataset
 
+from src.retrievers.fais_retriever import FAISRetriever
+from src.utils.log import get_logger
+from src.evaluation import evaluate
+from typing import cast
 
 logger = get_logger()
 
 
 if __name__ == '__main__':
+    dataset_name = "GroNLP/ik-nlp-22_slp"
+    paragraphs = load_dataset(dataset_name, "paragraphs")
+    questions = cast(DatasetDict, load_dataset(dataset_name, "questions"))
+
+    questions_test = questions["test"]
+
+    logger.info(questions)
+
     # Initialize retriever
     r = FAISRetriever()
 
-    # Retrieve example
-    scores, result = r.retrieve(
-        "What is the perplexity of a language model?")
+    # # Retrieve example
+    example_q = "What is the perplexity of a language model?"
+    scores, result = r.retrieve(example_q)
+
+    logger.info(
+        f"Example q: {example_q} answer: {result['text'][0]}")
 
     for i, score in enumerate(scores):
         logger.info(f"Result {i+1} (score: {score:.02f}):")
         logger.info(result['text'][i])
 
     # Compute overall performance
-    exact_match, f1_score = r.evaluate()
+    exact_match, f1_score = evaluate(
+        r, questions_test["question"], questions_test["answer"])
     logger.info(f"Exact match: {exact_match:.02f}\n"
                 f"F1-score: {f1_score:.02f}")

src/evaluation.py

@@ -1,4 +1,5 @@
-from typing import Callable, List
+from typing import Any, Callable, List
+from src.retrievers.base_retriever import Retriever
 
 from src.utils.string_utils import (lower, remove_articles, remove_punc,
                                     white_space_fix)
@@ -63,3 +64,29 @@ def f1(prediction: str, answer: str) -> float:
     rec = len(common_tokens) / len(answer_tokens)
 
     return 2 * (prec * rec) / (prec + rec)
+
+
+def evaluate(retriever: Retriever, questions: Any, answers: Any):
+    """Evaluates the entire model by computing F1-score and exact match on the
+    entire dataset.
+
+    Returns:
+        float: overall exact match
+        float: overall F1-score
+    """
+
+    predictions = []
+    scores = 0
+
+    # Currently just takes the first answer and does not look at scores yet
+    for question in questions:
+        score, result = retriever.retrieve(question, 1)
+        scores += score[0]
+        predictions.append(result['text'][0])
+
+    exact_matches = [exact_match(
+        predictions[i], answers[i]) for i in range(len(answers))]
+    f1_scores = [f1(
+        predictions[i], answers[i]) for i in range(len(answers))]
+
+    return sum(exact_matches) / len(exact_matches), sum(f1_scores) / len(f1_scores)

src/retrievers/base_retriever.py

@@ -0,0 +1,3 @@
+class Retriever():
+    def retrieve(self, query: str, k: int):
+        pass

src/{es_retriever.py → retrievers/es_retriever.py}

@@ -1,8 +1,10 @@
-class ESRetriever:
-    def __init__(self, dataset_name: str = "GroNLP/ik-nlp-22_slp"):
-        self.dataset_name = dataset_name
+from src.utils.log import get_logger
 
-    def _setup_data(self):
+logger = get_logger()
+
+
+class ESRetriever(Retriever):
+    def __init__(self, data_set):
         pass
 
     def retrieve(self, query: str, k: int):

src/{fais_retriever.py → retrievers/fais_retriever.py}

@@ -1,19 +1,27 @@
-# Hacky fix for FAISS error on macOS
-# See https://stackoverflow.com/a/63374568/4545692
 import os
 import os.path
 
 import torch
 from datasets import load_dataset
-from transformers import (DPRContextEncoder, DPRContextEncoderTokenizer,
-                          DPRQuestionEncoder, DPRQuestionEncoderTokenizer)
+from transformers import (
+    DPRContextEncoder,
+    DPRContextEncoderTokenizer,
+    DPRQuestionEncoder,
+    DPRQuestionEncoderTokenizer,
+)
 
-from src.evaluation import exact_match, f1
+from src.retrievers.base_retriever import Retriever
+from src.utils.log import get_logger
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
+# Hacky fix for FAISS error on macOS
+# See https://stackoverflow.com/a/63374568/4545692
+
 
+logger = get_logger()
 
-class FAISRetriever:
+
+class FAISRetriever(Retriever):
     """A class used to retrieve relevant documents based on some query.
     based on https://huggingface.co/docs/datasets/faiss_es#faiss.
     """
@@ -65,7 +73,7 @@ class FAISRetriever:
         # Load dataset
         ds = load_dataset(dataset_name, name="paragraphs")[
             "train"]  # type: ignore
-        print(ds)
+        logger.info(ds)
 
         if os.path.exists(embedding_path):
             # If we already have FAISS embeddings, load them from disk
@@ -115,32 +123,3 @@ class FAISRetriever:
         )
 
         return scores, results
-
-    def evaluate(self):
-        """Evaluates the entire model by computing F1-score and exact match on the
-        entire dataset.
-
-        Returns:
-            float: overall exact match
-            float: overall F1-score
-        """
-        questions_ds = load_dataset(
-            self.dataset_name, name="questions")['test']
-        questions = questions_ds['question']
-        answers = questions_ds['answer']
-
-        predictions = []
-        scores = 0
-
-        # Currently just takes the first answer and does not look at scores yet
-        for question in questions:
-            score, result = self.retrieve(question, 1)
-            scores += score[0]
-            predictions.append(result['text'][0])
-
-        exact_matches = [exact_match(
-            predictions[i], answers[i]) for i in range(len(answers))]
-        f1_scores = [f1(
-            predictions[i], answers[i]) for i in range(len(answers))]
-
-        return sum(exact_matches) / len(exact_matches), sum(f1_scores) / len(f1_scores)