Add longformer

.gitattributes

@@ -28,3 +28,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+src/models/dpr.faiss filter=lfs diff=lfs merge=lfs -text
+src/models/longformer.faiss filter=lfs diff=lfs merge=lfs -text

README.md

@@ -75,7 +75,10 @@ By default, the best answer along with its location in the book will be
 returned. If you want to generate more answers (say, a top-5), you can supply
 the `--top=5` option. The default retriever uses [FAISS](https://faiss.ai/), but
 you can also use [ElasticSearch](https://www.elastic.co/elastic-stack/) using
-the `--retriever=es` option.
+the `--retriever=es` option. You can also pick a language model using the
+`--lm` option, which accepts either `dpr` (Dense Passage Retrieval) or
+`longformer`. The language model is used to generate embeddings for FAISS, and
+is used to generate the answer.
 
 ### CLI overview
 
@@ -83,7 +86,7 @@ To get an overview of all available options, run `python query.py --help`. The
 options are also printed below.
 
 ```sh
-usage: query.py [-h] [--top int] [--retriever {faiss,es}] str
+usage: query.py [-h] [--top int] [--retriever {faiss,es}] [--lm {dpr,longformer}] str
 
 positional arguments:
   str                   The question to feed to the QA system
@@ -93,6 +96,8 @@ options:
   --top int, -t int     The number of answers to retrieve
   --retriever {faiss,es}, -r {faiss,es}
                         The retrieval method to use
+  --lm {dpr,longformer}, -l {dpr,longformer}
+                        The language model to use for the FAISS retriever
 ```
 
 

query.py

@@ -2,21 +2,48 @@ import argparse
 import torch
 import transformers
 
-from typing import List, Literal, Union, cast
+from typing import Dict, List, Literal, Tuple, cast
 from datasets import load_dataset, DatasetDict
 from dotenv import load_dotenv
 
+from src.readers.base_reader import Reader
+from src.readers.longformer_reader import LongformerReader
 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.retrievers.faiss_retriever import (
+    FaissRetriever,
+    FaissRetrieverOptions
+)
 from src.utils.preprocessing import context_to_reader_input
 from src.utils.log import get_logger
 
 
-def get_retriever(r: Union[Literal["es"], Literal["fais"]], paragraphs: DatasetDict) -> Retriever:
-    retriever = ESRetriever if r == "es" else FaissRetriever
-    return retriever(paragraphs)
+def get_retriever(paragraphs: DatasetDict,
+                  r: Literal["es", "faiss"],
+                  lm: Literal["dpr", "longformer"]) -> Retriever:
+    match (r, lm):
+        case "es", _:
+            return ESRetriever()
+        case "faiss", "dpr":
+            options = FaissRetrieverOptions.dpr("./src/models/dpr.faiss")
+            return FaissRetriever(paragraphs, options)
+        case "faiss", "longformer":
+            options = FaissRetrieverOptions.longformer(
+                "./src/models/longformer.faiss")
+            return FaissRetriever(paragraphs, options)
+        case _:
+            raise ValueError("Retriever options not recognized")
+
+
+def get_reader(lm: Literal["dpr", "longformer"]) -> Reader:
+    match lm:
+        case "dpr":
+            return DprReader()
+        case "longformer":
+            return LongformerReader()
+        case _:
+            raise ValueError("Language model not recognized")
 
 
 def print_name(contexts: dict, section: str, id: int):
@@ -51,7 +78,11 @@ def print_answers(answers: List[tuple], scores: List[float], contexts: dict):
         print()
 
 
-def probe(query: str, retriever: Retriever, reader: DprReader, num_answers: int = 5):
+def probe(query: str,
+          retriever: Retriever,
+          reader: Reader,
+          num_answers: int = 5) \
+          -> Tuple[List[tuple], List[float], Dict[str, List[str]]]:
     scores, contexts = retriever.retrieve(query)
     reader_input = context_to_reader_input(contexts)
     answers = reader.read(query, reader_input, num_answers)
@@ -63,7 +94,7 @@ def default_probe(query: str):
     # default probe is a probe that prints 5 answers with faiss
     paragraphs = cast(DatasetDict, load_dataset(
         "GroNLP/ik-nlp-22_slp", "paragraphs"))
-    retriever = get_retriever("faiss", paragraphs)
+    retriever = get_retriever(paragraphs, "faiss", "dpr")
     reader = DprReader()
 
     return probe(query, retriever, reader)
@@ -75,13 +106,20 @@ def main(args: argparse.Namespace):
         "GroNLP/ik-nlp-22_slp", "paragraphs"))
 
     # Retrieve
-    retriever = get_retriever(args.retriever, paragraphs)
-    reader = DprReader()
+    retriever = get_retriever(paragraphs, args.retriever, args.lm)
+    reader = get_reader(args.lm)
     answers, scores, contexts = probe(
-        args.query, retriever, reader, args.num_answers)
+        args.query, retriever, reader, args.top)
 
     # Print output
-    print_answers(answers, scores, contexts)
+    print("Question: " + args.query)
+    print("Answer(s):")
+    if args.lm == "dpr":
+        print_answers(answers, scores, contexts)
+    else:
+        answers = filter(lambda a: len(a[0].strip()) > 0, answers)
+        for pos, answer in enumerate(answers, start=1):
+            print(f"    - {answer[0].strip()}")
 
 
 if __name__ == "__main__":
@@ -94,13 +132,18 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         formatter_class=argparse.MetavarTypeHelpFormatter
     )
-    parser.add_argument("query", type=str,
-                        help="The question to feed to the QA system")
-    parser.add_argument("--top", "-t", type=int, default=1,
-                        help="The number of answers to retrieve")
-    parser.add_argument("--retriever", "-r", type=str.lower,
-                        choices=["faiss", "es"], default="faiss",
-                        help="The retrieval method to use")
+    parser.add_argument(
+        "query", type=str, help="The question to feed to the QA system")
+    parser.add_argument(
+        "--top", "-t", type=int, default=1,
+        help="The number of answers to retrieve")
+    parser.add_argument(
+        "--retriever", "-r", type=str.lower, choices=["faiss", "es"],
+        default="faiss", help="The retrieval method to use")
+    parser.add_argument(
+        "--lm", "-l", type=str.lower,
+        choices=["dpr", "longformer"], default="dpr",
+        help="The language model to use for the FAISS retriever")
 
     args = parser.parse_args()
     main(args)

src/models/{paragraphs_embedding.faiss → dpr.faiss}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fddf97865d5b1a967df90b7e2808bd27510cce633d55ed2af8328619828b168
+oid sha256:6bc0e5c38ddeb0a6a4daaf3ae98cd3e564f22ff9a263bc8867d0b363e828ccce
 size 5213229

src/models/longformer.faiss

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56b2616392540f4d2d8fa34d313a59c41572dca3ef5a683c7a8dbd2691418ea6
+size 5213229

src/readers/base_reader.py

@@ -0,0 +1,9 @@
+from typing import Dict, List, Tuple
+
+
+class Reader():
+    def read(self,
+             query: str,
+             context: Dict[str, List[str]],
+             num_answers: int) -> List[Tuple]:
+        raise NotImplementedError()

src/readers/dpr_reader.py

@@ -1,8 +1,10 @@
 from transformers import DPRReader, DPRReaderTokenizer
 from typing import List, Dict, Tuple
 
+from src.readers.base_reader import Reader
 
-class DprReader():
+
+class DprReader(Reader):
     def __init__(self) -> None:
         self._tokenizer = DPRReaderTokenizer.from_pretrained(
             "facebook/dpr-reader-single-nq-base")

src/readers/longformer_reader.py

@@ -0,0 +1,41 @@
+import torch
+from transformers import (
+    LongformerTokenizerFast,
+    LongformerForQuestionAnswering
+)
+from typing import List, Dict, Tuple
+
+from src.readers.base_reader import Reader
+
+
+class LongformerReader(Reader):
+    def __init__(self) -> None:
+        checkpoint = "valhalla/longformer-base-4096-finetuned-squadv1"
+        self.tokenizer = LongformerTokenizerFast.from_pretrained(checkpoint)
+        self.model = LongformerForQuestionAnswering.from_pretrained(checkpoint)
+
+    def read(self,
+             query: str,
+             context: Dict[str, List[str]],
+             num_answers=5) -> List[Tuple]:
+        answers = []
+
+        for text in context['texts']:
+            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)
+
+            start_logits = outputs.start_logits
+            end_logits = outputs.end_logits
+            all_tokens = self.tokenizer.convert_ids_to_tokens(
+                input_ids[0].tolist())
+            answer_tokens = all_tokens[
+                torch.argmax(start_logits):torch.argmax(end_logits) + 1]
+            answer = self.tokenizer.decode(
+                self.tokenizer.convert_tokens_to_ids(answer_tokens)
+            )
+            answers.append([answer, [], []])
+
+        return answers

src/retrievers/faiss_retriever.py

@@ -1,14 +1,19 @@
 import os
 import os.path
-
 import torch
-from datasets import DatasetDict, load_dataset
+
+from datasets import DatasetDict
+from dataclasses import dataclass
 from transformers import (
     DPRContextEncoder,
-    DPRContextEncoderTokenizer,
+    DPRContextEncoderTokenizerFast,
     DPRQuestionEncoder,
-    DPRQuestionEncoderTokenizer,
+    DPRQuestionEncoderTokenizerFast,
+    LongformerModel,
+    LongformerTokenizerFast
 )
+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
@@ -23,35 +28,99 @@ os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
 logger = get_logger()
 
 
+@dataclass
+class FaissRetrieverOptions:
+    ctx_encoder: PreTrainedModel
+    ctx_tokenizer: PreTrainedTokenizerFast
+    q_encoder: PreTrainedModel
+    q_tokenizer: PreTrainedTokenizerFast
+    embedding_path: str
+    lm: str
+
+    @staticmethod
+    def dpr(embedding_path: str):
+        return FaissRetrieverOptions(
+            ctx_encoder=DPRContextEncoder.from_pretrained(
+                "facebook/dpr-ctx_encoder-single-nq-base"
+            ),
+            ctx_tokenizer=DPRContextEncoderTokenizerFast.from_pretrained(
+                "facebook/dpr-ctx_encoder-single-nq-base"
+            ),
+            q_encoder=DPRQuestionEncoder.from_pretrained(
+                "facebook/dpr-question_encoder-single-nq-base"
+            ),
+            q_tokenizer=DPRQuestionEncoderTokenizerFast.from_pretrained(
+                "facebook/dpr-question_encoder-single-nq-base"
+            ),
+            embedding_path=embedding_path,
+            lm="dpr"
+        )
+
+    @staticmethod
+    def longformer(embedding_path: str):
+        encoder = LongformerModel.from_pretrained(
+            "allenai/longformer-base-4096"
+        )
+        tokenizer = LongformerTokenizerFast.from_pretrained(
+            "allenai/longformer-base-4096"
+        )
+        return FaissRetrieverOptions(
+            ctx_encoder=encoder,
+            ctx_tokenizer=tokenizer,
+            q_encoder=encoder,
+            q_tokenizer=tokenizer,
+            embedding_path=embedding_path,
+            lm="longformer"
+        )
+
+
 class FaissRetriever(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, paragraphs: DatasetDict, embedding_path: str = "./src/models/paragraphs_embedding.faiss") -> None:
+    def __init__(self, paragraphs: DatasetDict,
+                 options: FaissRetrieverOptions) -> None:
         torch.set_grad_enabled(False)
 
+        self.lm = options.lm
+
         # 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"
-        )
+        self.ctx_encoder = options.ctx_encoder
+        self.ctx_tokenizer = options.ctx_tokenizer
 
         # 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"
-        )
+        self.q_encoder = options.q_encoder
+        self.q_tokenizer = options.q_tokenizer
 
         self.paragraphs = paragraphs
-        self.embedding_path = embedding_path
+        self.embedding_path = options.embedding_path
 
         self.index = self._init_index()
 
+    def _embed_question(self, q):
+        match self.lm:
+            case "dpr":
+                tok = self.q_tokenizer(q, return_tensors="pt", truncation=True)
+                return self.q_encoder(**tok)[0][0].numpy()
+            case "longformer":
+                tok = self.q_tokenizer(q, return_tensors="pt")
+                return self.q_encoder(**tok).last_hidden_state[0][0].numpy()
+
+    def _embed_context(self, row):
+        p = row["text"]
+
+        match self.lm:
+            case "dpr":
+                tok = self.ctx_tokenizer(
+                    p, return_tensors="pt", truncation=True)
+                enc = self.ctx_encoder(**tok)[0][0].numpy()
+                return {"embeddings": enc}
+            case "longformer":
+                tok = self.ctx_tokenizer(p, return_tensors="pt")
+                enc = self.ctx_encoder(**tok).last_hidden_state[0][0].numpy()
+                return {"embeddings": enc}
+
     def _init_index(
             self,
             force_new_embedding: bool = False):
@@ -64,16 +133,8 @@ class FaissRetriever(Retriever):
                 'embeddings', self.embedding_path)  # type: ignore
             return ds
         else:
-            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
-            index = ds.map(embed)  # type: ignore
+            index = ds.map(self._embed_context)  # type: ignore
 
             index.add_faiss_index(column="embeddings")
 
@@ -86,12 +147,7 @@ class FaissRetriever(Retriever):
 
     @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)
-            return self.q_encoder(**tok)[0][0].numpy()
-
-        question_embedding = embed(query)
+        question_embedding = self._embed_question(query)
         scores, results = self.index.get_nearest_examples(
             "embeddings", question_embedding, k=k
         )