Add retriever based on DPR (WIP)

.gitignore

@@ -1,3 +1,6 @@
+# VS Code stuff
+.vscode/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -1,14 +1,27 @@
 # nlp-flashcard-project
 
-
 ## Todo voor progress meeting
 
-- Data inlezen/Repo klaarmaken
-- Proof of concept met UnifiedQA
-- Standaard QA model met de dataset
-- Papers verzamelen/lezen
-- Eerder werk bekijken, inspiratie opdoen voor research richting
+- [ ] Data inlezen/Repo klaarmaken
+- [ ] Proof of concept met UnifiedQA
+- [ ] Standaard QA model met de dataset
+- [ ] Papers verzamelen/lezen
+- [ ] Eerder werk bekijken, inspiratie opdoen voor research richting
+
+## Overview
+
+De meeste QA systemen bestaan uit twee onderdelen:
+
+- Een retriever. Die haalt adhv de vraag _k_ relevante stukken context op, bv.
+  met `tf-idf`.
+- Een model dat het antwoord genereert. Wat je hier precies gebruikt hangt af
+  van de manier van question answering:
+  - Voor **extractive QA** gebruik je een reader;
+  - Voor **generative QA** gebruik je een generator.
+
+  Beide werken op basis van een language model.
 
 ## Handige info
 
-Huggingface QA tutorial: <https://huggingface.co/docs/transformers/tasks/question_answering#finetune-with-tensorflow>
+- Huggingface QA tutorial: <https://huggingface.co/docs/transformers/tasks/question_answering#finetune-with-tensorflow>
+- Overview van open-domain question answering technieken: <https://lilianweng.github.io/posts/2020-10-29-odqa/>

base_model/reader.py

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

base_model/retriever.py

@@ -0,0 +1,86 @@
+from transformers import DPRContextEncoder, DPRContextEncoderTokenizer, \
+                         DPRQuestionEncoder, DPRQuestionEncoderTokenizer
+from datasets import load_dataset
+import torch
+
+
+class Retriever():
+    """A class used to retrieve relevant documents based on some query.
+    based on https://huggingface.co/docs/datasets/faiss_es#faiss.
+    """
+
+    def __init__(self, dataset: str = "GroNLP/ik-nlp-22_slp") -> None:
+        """Initialize the retriever
+
+        Args:
+            dataset (str, optional): The dataset to train on. Assumes the
+            information is stored in a column named 'text'. Defaults to
+            "GroNLP/ik-nlp-22_slp".
+        """
+        torch.set_grad_enabled(False)
+
+        # Context encoding and tokenization
+        self.ctx_encoder = DPRContextEncoder.from_pretrained(
+            "facebook/dpr-ctx_encoder-single-nq-base")
+        self.ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained(
+            "facebook/dpr-ctx_encoder-single-nq-base")
+
+        # Question encoding and tokenization
+        self.q_encoder = DPRQuestionEncoder.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base")
+        self.q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base")
+
+        # Dataset building
+        self.dataset = self.__init_dataset(dataset)
+
+    def __init_dataset(self, dataset: str):
+        """Loads the dataset and adds FAISS embeddings.
+
+        Args:
+            dataset (str): A HuggingFace dataset name.
+
+        Returns:
+            Dataset: A dataset with a new column 'embeddings' containing FAISS
+            embeddings.
+        """
+        # TODO: save ds w/ embeddings to disk and retrieve it if it already exists
+
+        # Load dataset
+        ds = load_dataset(dataset, name='paragraphs')['train']
+
+        def embed(row):
+            # Inline helper function to perform embedding
+            p = row['text']
+            tok = self.ctx_tokenizer(p, return_tensors='pt', truncation=True)
+            enc = self.ctx_encoder(**tok)[0][0].numpy()
+            return {'embeddings': enc}
+
+        # Add FAISS embeddings
+        ds_with_embeddings = ds.map(embed)
+
+        # Todo: this throws a weird error.
+        ds_with_embeddings.add_faiss_index(column='embeddings')
+        return ds_with_embeddings
+
+    def retrieve(self, query: str, k: int = 5):
+        """Retrieve the top k matches for a search query.
+
+        Args:
+            query (str): A search query
+            k (int, optional): The number of documents to retrieve. Defaults to
+            5.
+
+        Returns:
+            tuple: A tuple of lists of scores and results.
+        """
+
+        def embed(q):
+            # Inline helper function to perform embedding
+            tok = self.q_tokenizer(q, return_tensors='pt', truncation=True)
+            return self.q_encoder(**tok)[0][0].numpy()
+
+        question_embedding = embed(query)
+        scores, results = self.dataset.get_nearest_examples(
+            'embeddings', question_embedding, k=k)
+        return scores, results

main.py

@@ -0,0 +1,14 @@
+from base_model.retriever import Retriever
+
+if __name__ == '__main__':
+    # Initialize retriever
+    r = Retriever()
+
+    # Retrieve example
+    retrieved = r.retrieve(
+        "When is a stochastic process said to be stationary?")
+
+    for i, (score, result) in enumerate(retrieved):
+        print(f"Result {i+1} (score: {score * 100:.02f}:")
+        print(result['text'][0])
+        print()  # Newline

pyproject.toml

@@ -7,8 +7,13 @@ authors = ["Your Name <you@example.com>"]
 [tool.poetry.dependencies]
 python = "^3.8"
 numpy = "^1.22.3"
+transformers = "^4.17.0"
+torch = "^1.11.0"
+datasets = "^1.18.4"
+faiss-cpu = "^1.7.2"
 
 [tool.poetry.dev-dependencies]
+flake8 = "^4.0.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]