Add reader

main.py

@@ -1,12 +1,21 @@
 from datasets import DatasetDict, load_dataset
 
-from src.retrievers.fais_retriever import FAISRetriever
+from src.readers.dpr_reader import DprReader
+from src.retrievers.faiss_retriever import FaissRetriever
 from src.utils.log import get_logger
-from src.evaluation import evaluate
+# from src.evaluation import evaluate
 from typing import cast
 
-logger = get_logger()
+from src.utils.preprocessing import result_to_reader_input
+
+import torch
+import transformers
+import os
 
+os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = '1'
+
+logger = get_logger()
+transformers.logging.set_verbosity_error()
 
 if __name__ == '__main__':
     dataset_name = "GroNLP/ik-nlp-22_slp"
@@ -15,24 +24,44 @@ if __name__ == '__main__':
 
     questions_test = questions["test"]
 
-    logger.info(questions)
+    # logger.info(questions)
 
     # Initialize retriever
-    r = FAISRetriever()
+    retriever = FaissRetriever()
+
+    # Retrieve example
+    example_q = questions_test.shuffle()["question"][0]
+    scores, result = retriever.retrieve(example_q)
+
+    reader_input = result_to_reader_input(result)
+
+    # 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]))
 
-    # # Retrieve example
-    example_q = "What is the perplexity of a language model?"
-    scores, result = r.retrieve(example_q)
+    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
 
-    logger.info(
-        f"Example q: {example_q} answer: {result['text'][0]}")
+    # print(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])
+    # for i, score in enumerate(scores):
+    #     print(f"Result {i+1} (score: {score:.02f}):")
+    #     print(result['text'][i])
 
-    # Compute overall performance
-    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}")
+    # # Compute overall performance
+    # exact_match, f1_score = evaluate(
+    #     r, questions_test["question"], questions_test["answer"])
+    # print(f"Exact match: {exact_match:.02f}\n", f"F1-score: {f1_score:.02f}")

src/readers/dpr_reader.py

@@ -0,0 +1,27 @@
+from transformers import DPRReader, DPRReaderTokenizer
+from typing import List, Dict, Tuple
+
+
+class DprReader():
+    def __init__(self) -> None:
+        self._tokenizer = DPRReaderTokenizer.from_pretrained(
+            "facebook/dpr-reader-single-nq-base")
+        self._model = DPRReader.from_pretrained(
+            "facebook/dpr-reader-single-nq-base"
+        )
+
+    def read(self, query: str, context: Dict[str, List[str]]) -> List[Tuple]:
+        encoded_inputs = self._tokenizer(
+            questions=query,
+            titles=context['titles'],
+            texts=context['texts'],
+            return_tensors='pt',
+            truncation=True,
+            padding=True
+        )
+        outputs = self._model(**encoded_inputs)
+
+        predicted_spans = self._tokenizer.decode_best_spans(
+            encoded_inputs, outputs)
+
+        return predicted_spans

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

@@ -13,15 +13,15 @@ from transformers import (
 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
+os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
 
 
 logger = get_logger()
 
 
-class FAISRetriever(Retriever):
+class FaissRetriever(Retriever):
     """A class used to retrieve relevant documents based on some query.
     based on https://huggingface.co/docs/datasets/faiss_es#faiss.
     """
@@ -56,14 +56,16 @@ class FAISRetriever(Retriever):
         self.dataset_name = dataset_name
         self.dataset = self._init_dataset(dataset_name)
 
-    def _init_dataset(self,
-                      dataset_name: str,
-                      embedding_path: str = "./models/paragraphs_embedding.faiss"):
+    def _init_dataset(
+            self,
+            dataset_name: str,
+            embedding_path: str = "./src/models/paragraphs_embedding.faiss",
+            force_new_embedding: bool = False):
         """Loads the dataset and adds FAISS embeddings.
 
         Args:
             dataset (str): A HuggingFace dataset name.
-            fname (str): The name to use to save the embeddings to disk for 
+            fname (str): The name to use to save the embeddings to disk for
             faster loading after the first run.
 
         Returns:
@@ -73,9 +75,8 @@ class FAISRetriever(Retriever):
         # Load dataset
         ds = load_dataset(dataset_name, name="paragraphs")[
             "train"]  # type: ignore
-        logger.info(ds)
 
-        if os.path.exists(embedding_path):
+        if not force_new_embedding and os.path.exists(embedding_path):
             # If we already have FAISS embeddings, load them from disk
             ds.load_faiss_index('embeddings', embedding_path)  # type: ignore
             return ds
@@ -95,7 +96,7 @@ class FAISRetriever(Retriever):
             ds_with_embeddings.add_faiss_index(column="embeddings")
 
             # save dataset w/ embeddings
-            os.makedirs("./models/", exist_ok=True)
+            os.makedirs("./src/models/", exist_ok=True)
             ds_with_embeddings.save_faiss_index("embeddings", embedding_path)
 
             return ds_with_embeddings

src/utils/preprocessing.py

@@ -0,0 +1,35 @@
+from typing import Dict, List
+
+
+def result_to_reader_input(result: Dict[str, List[str]]) \
+        -> Dict[str, List[str]]:
+    """Takes the output of the retriever and turns it into a format the reader
+    understands.
+
+    Args:
+        result (Dict[str, List[str]]): The result from the retriever
+    """
+
+    # Take the number of valeus of an arbitrary item as the number of entries
+    # (This should always be valid)
+    num_entries = len(result['n_chapter'])
+
+    # Prepare result
+    reader_result = {
+        'titles': [],
+        'texts': []
+    }
+
+    for n in range(num_entries):
+        # Get the most specific title
+        if result['subsection'][n] != 'nan':
+            title = result['subsection'][n]
+        elif result['section'][n] != 'nan':
+            title = result['section'][n]
+        else:
+            title = result['chapter'][n]
+
+        reader_result['titles'].append(title)
+        reader_result['texts'].append(result['text'][n])
+
+    return reader_result