Upload kilt_create_dpr_support_docs.py

kilt_create_dpr_support_docs.py

@@ -0,0 +1,109 @@
+import argparse
+import json
+import os
+
+import faiss
+import torch
+from datasets import load_dataset, Dataset
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, DPRQuestionEncoder, DPRContextEncoder
+
+from common import articles_to_paragraphs, embed_questions, embed_passages, create_kilt_datapoint, \
+    kilt_wikipedia_columns
+from common import kilt_wikipedia_paragraph_columns as columns
+
+
+def generate_support_docs(args):
+    dims = 128
+    min_chars_per_passage = 200
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+    lfqa = load_dataset("vblagoje/lfqa")
+
+    ctx_tokenizer = AutoTokenizer.from_pretrained(args.ctx_encoder_name)
+    ctx_model = DPRContextEncoder.from_pretrained(args.ctx_encoder_name).to(device)
+    _ = ctx_model.eval()
+
+    question_tokenizer = AutoTokenizer.from_pretrained(args.question_encoder_name)
+    question_model = DPRQuestionEncoder.from_pretrained(args.question_encoder_name).to(device)
+    _ = question_model.eval()
+
+    kilt_wikipedia = load_dataset("kilt_wikipedia", split="full")
+
+    kilt_wikipedia_paragraphs = kilt_wikipedia.map(articles_to_paragraphs, batched=True,
+                                                   remove_columns=kilt_wikipedia_columns,
+                                                   batch_size=512,
+                                                   cache_file_name=f"../data/wiki_kilt_paragraphs_full.arrow",
+                                                   desc="Expanding wiki articles into paragraphs")
+
+    # use paragraphs that are not simple fragments or very short sentences
+    # Wikipedia Faiss index needs to fit into a 16 Gb GPU
+    kilt_wikipedia_paragraphs = kilt_wikipedia_paragraphs.filter(
+        lambda x: (x["end_character"] - x["start_character"]) > min_chars_per_passage)
+
+    def query_index(question, topk=7):
+        topk = topk * 3  # grab 3x results and filter for word count
+        question_embedding = embed_questions(question_model, question_tokenizer, [question])
+        scores, wiki_passages = kilt_wikipedia_paragraphs.get_nearest_examples("embeddings", question_embedding, k=topk)
+
+        retrieved_examples = []
+        r = list(zip(wiki_passages[k] for k in columns))
+        for i in range(topk):
+            retrieved_examples.append({k: v for k, v in zip(columns, [r[j][0][i] for j in range(len(columns))])})
+
+        return retrieved_examples
+
+    def create_support_doc(dataset: Dataset, output_filename: str):
+        progress_bar = tqdm(range(len(dataset)), desc="Creating supporting docs")
+
+        with open(output_filename, "w") as fp:
+            for example in dataset:
+                wiki_passages = query_index(example["title"])
+                kilt_dp = create_kilt_datapoint(example, columns, wiki_passages)
+                json.dump(kilt_dp, fp)
+                fp.write("\n")
+                progress_bar.update(1)
+
+    if not os.path.isfile(args.index_file_name):
+        def embed_passages_for_retrieval(examples):
+            return embed_passages(ctx_model, ctx_tokenizer, examples, max_length=128)
+
+        paragraphs_embeddings = kilt_wikipedia_paragraphs.map(embed_passages_for_retrieval,
+                                                              batched=True, batch_size=512,
+                                                              cache_file_name=args.encoded_kilt_file_name,
+                                                              desc="Creating faiss index")
+
+        paragraphs_embeddings.add_faiss_index(column="embeddings", custom_index=faiss.IndexFlatIP(dims))
+        paragraphs_embeddings.save_faiss_index("embeddings", args.index_file_name)
+
+    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", args.index_file_name, device=0)
+    create_support_doc(lfqa["train"], "lfqa_dpr_train_precomputed_dense_docs.json")
+    create_support_doc(lfqa["validation"], "lfqa_dpr_validation_precomputed_dense_docs.json")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Creates support docs for seq2seq model training")
+    parser.add_argument(
+        "--ctx_encoder_name",
+        default="vblagoje/dpr-ctx_encoder-single-lfqa-base",
+        help="Question encoder to use",
+    )
+    parser.add_argument(
+        "--question_encoder_name",
+        default="vblagoje/dpr-question_encoder-single-lfqa-base",
+        help="Question encoder to use",
+    )
+
+    parser.add_argument(
+        "--index_file_name",
+        default="../data/kilt_dpr_wikipedia_first.faiss",
+        help="Faiss index with passage embeddings",
+    )
+
+    parser.add_argument(
+        "--encoded_kilt_file_name",
+        default="../data/kilt_embedded.arrow",
+        help="Encoded KILT file name",
+    )
+
+    main_args, _ = parser.parse_known_args()
+    generate_support_docs(main_args)