Add files via upload

util/common.py

@@ -0,0 +1,120 @@
+import re
+
+import torch
+
+kilt_wikipedia_columns = ['kilt_id', 'wikipedia_id', 'wikipedia_title', 'text', 'anchors', 'categories',
+                          'wikidata_info', 'history']
+
+kilt_wikipedia_paragraph_columns = ['wikipedia_id', 'start_paragraph_id', 'start_character', 'end_paragraph_id',
+                                    'end_character', 'title', 'section', 'text']
+
+
+def clean_question(text):
+    result = cleanup_references(text)
+    result = result.replace("\n", " ")
+    result = re.sub(r"\s\s+", " ", result)
+    result = result.replace("[deleted]", "")
+    return result.lower().strip()
+
+
+def cleanup_references(text):
+    # URL reference where we need to remove both the link text and URL
+    # ...and this letter is used by most biographers as the cornerstone of Lee's personal
+    # views on slavery ([1](_URL_2_ & pg=PA173), [2](_URL_1_), [3](_URL_5_)).
+    # ...and this letter is used by most biographers as the cornerstone of Lee's personal views on slavery.
+    result = re.sub(r"[\(\s]*\[\d+\]\([^)]+\)[,)]*", "", text, 0, re.MULTILINE)
+
+    # URL reference where we need to preserve link text but remove URL
+    # At the outbreak of the Civil War, [Leyburn left his church](_URL_19_) and joined the South.
+    # At the outbreak of the Civil War, Leyburn left his church and joined the South.
+    result = re.sub(r"\[([^]]+)\]\([^)]+\)", "\\1", result, 0, re.MULTILINE)
+
+    # lastly remove just dangling _URL_[0-9]_ URL references
+    result = re.sub(r"_URL_\d_", "", result, 0, re.MULTILINE)
+    return result
+
+
+def clean_answer(text):
+    result = cleanup_references(text)
+    result = result.replace("\n", " ")
+    result = re.sub(r"\s\s+", " ", result)
+    result = re.sub(r"BULLET::::-", "", result)
+    return trim(result.strip())
+
+
+def trim(text, word_count: int = 100):
+    return " ".join(text.split(" ")[:word_count])
+
+
+def articles_to_paragraphs(examples):
+    ids, titles, sections, texts, start_ps, end_ps, start_cs, end_cs = [], [], [], [], [], [], [], []
+    for bidx, example in enumerate(examples["text"]):
+        last_section = ""
+        for idx, p in enumerate(example["paragraph"]):
+            if "Section::::" in p:
+                last_section = p
+            ids.append(examples["wikipedia_id"][bidx])
+            titles.append(examples["wikipedia_title"][bidx])
+            sections.append(last_section)
+            texts.append(p)
+            start_ps.append(idx)
+            end_ps.append(idx)
+            start_cs.append(0)
+            end_cs.append(len(p))
+
+    return {"wikipedia_id": ids, "title": titles,
+            "section": sections, "text": texts,
+            "start_paragraph_id": start_ps, "end_paragraph_id": end_ps,
+            "start_character": start_cs,
+            "end_character": end_cs
+            }
+
+
+def create_kilt_datapoint(eli5_example, columns, wiki_passages, min_length=20, topk=7):
+    res_list = [dict([(k, p[k]) for k in columns]) for p in wiki_passages]
+    res_list = [res for res in res_list if len(res["text"].split()) > min_length][:topk]
+
+    # make a KILT data point
+    # see https://github.com/facebookresearch/KILT#kilt-data-format
+    output = []
+    for a in eli5_example["answers"]["text"]:
+        output.append({"answer": a})
+
+    output.append({"provenance": [
+        # evidence set for the answer from the KILT ks
+        {
+            "wikipedia_id": r["wikipedia_id"],  # *mandatory*
+            "title": r["title"],
+            "section": r["section"],
+            "start_paragraph_id": r["start_paragraph_id"],
+            "start_character": r["start_character"],
+            "end_paragraph_id": r["end_paragraph_id"],
+            "end_character": r["end_character"],
+            "text": r["text"],
+            "bleu_score": None,  # wrt original evidence
+            "meta": None  # dataset/task specific
+        } for r in res_list
+    ]})
+    return {"id": eli5_example["q_id"],
+            "input": eli5_example["title"],
+            "output": output,  # each element is an answer or provenance (can have multiple of each)
+            "meta": None  # dataset/task specific
+            }
+
+
+def embed_questions(question_model, question_tokenizer, questions, max_length=128, device="cuda:0"):
+    query = question_tokenizer(questions, max_length=max_length, padding="max_length", truncation=True,
+                               return_tensors="pt")
+    with torch.no_grad():
+        q_reps = question_model(query["input_ids"].to(device),
+                                query["attention_mask"].to(device)).pooler_output
+    return q_reps.cpu().numpy()
+
+
+def embed_passages(ctx_model, ctx_tokenizer, passages, max_length=128, device="cuda:0"):
+    p = ctx_tokenizer(passages["text"], max_length=max_length, padding="max_length",
+                      truncation=True, return_tensors="pt")
+    with torch.no_grad():
+        a_reps = ctx_model(p["input_ids"].to(device),
+                           p["attention_mask"].to(device)).pooler_output
+    return {"embeddings": a_reps.cpu().numpy()}

util/create_dpr_training_from_dataset.py

@@ -0,0 +1,103 @@
+import argparse
+import random
+import json
+import re
+
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import semantic_search, cos_sim
+from tqdm.auto import tqdm
+from datasets import load_dataset
+
+from common import clean_answer, clean_question
+
+
+def find_hard_negative_ctxs(dataset, dataset_embeddings, embedding_index: int,
+                            exclude_answer_patterns, similarity_threshold=[0.5, 0.6], k=25, min_count=3):
+    hard_negative_ctxs = []
+    results = semantic_search(dataset_embeddings[embedding_index], dataset_embeddings, top_k=k,
+                              score_function=cos_sim)
+    # list if dicts
+    # [{'corpus_id': 8, 'score': -0.019427383318543434},
+    #  ...
+    # {'corpus_id': 10, 'score': -0.09040290117263794}]
+    # hard negative are most similar and negatives are most disimilar to embedding_index
+    hard_negative_results = results[0][1:k + 1]
+    assert len(hard_negative_results) > min_count * 2
+    for r in hard_negative_results:
+        example = dataset[r["corpus_id"]]
+        if similarity_threshold[0] < r["score"] <= similarity_threshold[1]:
+            for a in example["answers"]["text"]:
+                hard_negative_ctxs.append({"title": "", "text": clean_answer(a)})
+        if len(hard_negative_ctxs) > min_count:
+            break
+    return hard_negative_ctxs[:min_count]
+
+
+def find_negative_ctxs(dataset, dataset_embeddings, embedding_index: int,
+                       exclude_answer_patterns, similarity_threshold=0.1, k=7, min_count=3):
+    negative_ctxs = []
+    random_sample = random.sample(range(len(dataset_embeddings)), k * 20)
+    similarities = cos_sim(dataset_embeddings[embedding_index], dataset_embeddings[random_sample])[0].tolist()
+    for idx, score in enumerate(similarities):
+        if score < similarity_threshold:
+            example = dataset[random_sample[idx]]
+            for a in example["answers"]["text"]:
+                negative_ctxs.append({"title": "", "text": clean_answer(a)})
+        if len(negative_ctxs) > min_count:
+            break
+    return negative_ctxs[:min_count]
+
+
+def generate_dpr_training_file(args):
+    embedder = SentenceTransformer(args.embedding_model)
+
+    eli5_train_set = load_dataset("vblagoje/lfqa", split="train")
+    eli5_validation_set = load_dataset("vblagoje/lfqa", split="validation")
+    eli5_test_set = load_dataset("vblagoje/lfqa", split="test")
+
+    train_set = embedder.encode([example["title"] for example in eli5_train_set], convert_to_tensor=True,
+                                show_progress_bar=True)
+    validation_set = embedder.encode([example["title"] for example in eli5_validation_set], convert_to_tensor=True,
+                                     show_progress_bar=True)
+
+    test_set = embedder.encode([example["title"] for example in eli5_test_set], convert_to_tensor=True,
+                               show_progress_bar=True)
+    exclude_answer_patterns = [re.compile("not sure what you"), re.compile("\n\n >")]
+    for dataset_name, dataset, dataset_embeddings in zip(["train", "validation", "test"],
+                                                         [eli5_train_set, eli5_validation_set, eli5_test_set],
+                                                         [train_set, validation_set, test_set]):
+        min_elements = 3
+        skip_count = 0
+        progress_bar = tqdm(range(len(dataset)), desc="Creating DPR formatted question/passage docs")
+        with open('eli5-dpr-' + dataset_name + '.jsonl', 'w') as fp:
+            for idx, example in enumerate(dataset):
+                negative_ctxs = find_negative_ctxs(dataset, dataset_embeddings, idx, exclude_answer_patterns)
+                hard_negative_ctxs = find_hard_negative_ctxs(dataset, dataset_embeddings, idx, exclude_answer_patterns)
+                positive_context = [{"text": clean_answer(a), "title": ""} for a in example["answers"]["text"] if
+                                    not any([p.search(a) for p in exclude_answer_patterns])]
+                if not positive_context:
+                    positive_context = [{"text": clean_answer(a), "title": ""} for a in example["answers"]["text"]]
+                if len(positive_context) > 0 and len(negative_ctxs) > 0 and len(hard_negative_ctxs) >= min_elements:
+                    json.dump({"id": example["q_id"],
+                               "question": clean_question(example["title"]),
+                               "positive_ctxs": positive_context[:min_elements],
+                               "negative_ctxs": negative_ctxs[:min_elements],
+                               "hard_negative_ctxs": hard_negative_ctxs[:min_elements]}, fp)
+                    fp.write("\n")
+                else:
+                    skip_count += 1
+                progress_bar.update(1)
+
+        print(f"Skipped {skip_count} questions")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Creates DPR training file from LFQA dataset")
+    parser.add_argument(
+        "--embedding_model",
+        default="all-mpnet-base-v2",
+        help="Embedding model to use for question encoding and semantic search",
+    )
+
+    main_args, _ = parser.parse_known_args()
+    generate_dpr_training_file(main_args)

util/create_dpr_training_from_faiss.py

@@ -0,0 +1,144 @@
+import argparse
+import json
+
+import torch
+from datasets import load_dataset
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import DPRQuestionEncoder
+
+from common import embed_questions, clean_question, articles_to_paragraphs, kilt_wikipedia_columns
+from common import kilt_wikipedia_paragraph_columns as columns
+
+
+def generate_dpr_training_file(args):
+    n_negatives = 7
+    min_chars_per_passage = 200
+
+    def query_index(question, topk=(n_negatives * args.n_positives) * 2):
+        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 find_positive_and_hard_negative_ctxs(dataset_index: int, n_positive=1, device="cuda:0"):
+        positive_context_list = []
+        hard_negative_context_list = []
+        example = dataset[dataset_index]
+        question = clean_question(example['title'])
+        passages = query_index(question)
+        passages = [dict([(k, p[k]) for k in columns]) for p in passages]
+        q_passage_pairs = [[question, f"{p['title']} {p['text']}" if args.use_title else p["text"]] for p in passages]
+
+        features = ce_tokenizer(q_passage_pairs, padding="max_length", max_length=256, truncation=True,
+                                return_tensors="pt")
+        with torch.no_grad():
+            passage_scores = ce_model(features["input_ids"].to(device),
+                                      features["attention_mask"].to(device)).logits
+
+        for p_idx, p in enumerate(passages):
+            p["score"] = passage_scores[p_idx].item()
+
+        # order by scores
+        def score_passage(item):
+            return item["score"]
+
+        # pick the most relevant as the positive answer
+        best_passage_list = sorted(passages, key=score_passage, reverse=True)
+        for idx, item in enumerate(best_passage_list):
+            if idx < n_positive:
+                positive_context_list.append({"title": item["title"], "text": item["text"]})
+            else:
+                break
+
+        # least relevant as hard_negative
+        worst_passage_list = sorted(passages, key=score_passage, reverse=False)
+        for idx, hard_negative in enumerate(worst_passage_list):
+            if idx < n_negatives * n_positive:
+                hard_negative_context_list.append({"title": hard_negative["title"], "text": hard_negative["text"]})
+            else:
+                break
+        assert len(positive_context_list) * n_negatives == len(hard_negative_context_list)
+        return positive_context_list, hard_negative_context_list
+
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+
+    question_model = DPRQuestionEncoder.from_pretrained(args.question_encoder_name).to(device)
+    question_tokenizer = AutoTokenizer.from_pretrained(args.question_encoder_name)
+    _ = question_model.eval()
+
+    ce_model = AutoModelForSequenceClassification.from_pretrained('cross-encoder/ms-marco-MiniLM-L-4-v2').to(device)
+    ce_tokenizer = AutoTokenizer.from_pretrained('cross-encoder/ms-marco-MiniLM-L-4-v2')
+    _ = ce_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)
+
+    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", args.index_file_name, device=0)
+
+    eli5_train_set = load_dataset("vblagoje/lfqa", split="train")
+    eli5_validation_set = load_dataset("vblagoje/lfqa", split="validation")
+    eli5_test_set = load_dataset("vblagoje/lfqa", split="test")
+
+    for dataset_name, dataset in zip(["train", "validation", "test"], [eli5_train_set,
+                                                                       eli5_validation_set,
+                                                                       eli5_test_set]):
+
+        progress_bar = tqdm(range(len(dataset)), desc=f"Creating DPR formatted {dataset_name} file")
+        with open('eli5-dpr-' + dataset_name + '.jsonl', 'w') as fp:
+            for idx, example in enumerate(dataset):
+                negative_start_idx = 0
+                positive_context, hard_negative_ctxs = find_positive_and_hard_negative_ctxs(idx, args.n_positives,
+                                                                                            device)
+                for pc in positive_context:
+                    hnc = hard_negative_ctxs[negative_start_idx:negative_start_idx + n_negatives]
+                    json.dump({"id": example["q_id"],
+                               "question": clean_question(example["title"]),
+                               "positive_ctxs": [pc],
+                               "hard_negative_ctxs": hnc}, fp)
+                    fp.write("\n")
+                    negative_start_idx += n_negatives
+                progress_bar.update(1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Creates DPR training file")
+    parser.add_argument(
+        "--use_title",
+        action="store_true",
+        help="If true, use title in addition to passage text for passage embedding",
+    )
+    parser.add_argument(
+        "--n_positives",
+        default=3,
+        help="Number of positive samples per question",
+    )
+    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",
+    )
+
+    main_args, _ = parser.parse_known_args()
+    generate_dpr_training_file(main_args)

util/create_faiss_index.py

@@ -0,0 +1,67 @@
+import argparse
+import os
+
+import faiss
+import torch
+from datasets import load_dataset
+from transformers import AutoTokenizer, DPRContextEncoder
+
+from common import articles_to_paragraphs, embed_passages
+
+
+def create_faiss(args):
+    dims = 128
+    min_chars_per_passage = 200
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+
+    ctx_tokenizer = AutoTokenizer.from_pretrained(args.ctx_encoder_name)
+    ctx_model = DPRContextEncoder.from_pretrained(args.ctx_encoder_name).to(device)
+    _ = ctx_model.eval()
+
+    kilt_wikipedia = load_dataset("kilt_wikipedia", split="full")
+    kilt_wikipedia_columns = ['kilt_id', 'wikipedia_id', 'wikipedia_title', 'text', 'anchors', 'categories',
+                              'wikidata_info', 'history']
+
+    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)
+
+    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="../data/kilt_embedded.arrow",
+                                                              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)
+    else:
+        print(f"Faiss index already exists {args.index_file_name}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Creates Faiss Wikipedia index file")
+
+    parser.add_argument(
+        "--ctx_encoder_name",
+        default="vblagoje/dpr-ctx_encoder-single-lfqa-base",
+        help="Encoding model to use for passage encoding",
+    )
+
+    parser.add_argument(
+        "--index_file_name",
+        default="../data/kilt_dpr_wikipedia.faiss",
+        help="Faiss index file with passage embeddings",
+    )
+
+    main_args, _ = parser.parse_known_args()
+    create_faiss(main_args)

util/eval_generate.py

@@ -0,0 +1,140 @@
+import argparse
+import json
+import os
+
+import torch
+from datasets import load_dataset
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DPRQuestionEncoder
+
+from common import articles_to_paragraphs, kilt_wikipedia_columns
+from common import kilt_wikipedia_paragraph_columns as columns
+
+
+def eval_generate(args):
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+    question_tokenizer = AutoTokenizer.from_pretrained(args.question_encoder_name)
+    question_model = DPRQuestionEncoder.from_pretrained(args.question_encoder_name).to(device)
+    _ = question_model.eval()
+
+    eli5_tokenizer = AutoTokenizer.from_pretrained('vblagoje/bart_eli5')
+    eli5_model = AutoModelForSeq2SeqLM.from_pretrained('vblagoje/bart_eli5').to(device)
+    _ = eli5_model.eval()
+
+    min_snippet_length = 20
+    topk = 21
+    min_chars_per_passage = 200
+    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=256,
+                                                   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
+    kilt_wikipedia_paragraphs = kilt_wikipedia_paragraphs.filter(
+        lambda x: (x["end_character"] - x["start_character"]) > min_chars_per_passage)
+    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", args.index_file_name, device=0)
+
+    def embed_questions_for_retrieval(questions):
+        query = question_tokenizer(questions, max_length=128, padding=True, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            q_reps = question_model(query["input_ids"].to(device),
+                                    query["attention_mask"].to(device)).pooler_output
+        return q_reps.cpu().numpy()
+
+    def query_index(question):
+        question_embedding = embed_questions_for_retrieval([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_kilt_datapoint(q_id, query, answer, res_list):
+        # make a KILT data point
+        # see https://github.com/facebookresearch/KILT#kilt-data-format
+
+        provenance = [{
+            "wikipedia_id": r["wikipedia_id"],  # *mandatory*
+            "title": r["title"],
+            "section": r["section"],
+            "start_paragraph_id": r["start_paragraph_id"],
+            "start_character": r["start_character"],
+            "end_paragraph_id": r["end_paragraph_id"],
+            "end_character": r["end_character"],
+            "text": r["text"],
+            "bleu_score": None,  # wrt original evidence
+            "meta": None  # dataset/task specific
+        } for r in res_list]
+
+        output = [{"answer": answer, "provenance": provenance}]
+
+        return {"id": q_id,
+                "input": query,
+                "output": output,  # each element is an answer or provenance (can have multiple of each)
+                "meta": None  # dataset/task specific
+                }
+
+    kilt_output = []
+    with open(args.kilt_input_file, "r") as f:
+        kilt_items = [json.loads(x) for x in f.read().strip().split("\n")]
+        progress_bar = tqdm(range(len(kilt_items)), desc="Creating KILT response document")
+        for idx, item in enumerate(kilt_items):
+            query = item["input"]
+            res_list = query_index(query)
+
+            res_list = [res for res in res_list if len(res["text"].split()) > min_snippet_length][:int(topk / 3)]
+            documents = [res["text"] for res in res_list]
+            conditioned_doc = "<P> " + " <P> ".join([d for d in documents])
+
+            query_and_docs = "question: {} context: {}".format(query, conditioned_doc)
+
+            model_input = eli5_tokenizer(query_and_docs, truncation=True, padding=True, return_tensors="pt")
+            generated_answers_encoded = eli5_model.generate(input_ids=model_input["input_ids"].to(device),
+                                                            attention_mask=model_input["attention_mask"].to(device),
+                                                            min_length=50,
+                                                            max_length=250,
+                                                            do_sample=False,
+                                                            early_stopping=True,
+                                                            num_beams=8,
+                                                            temperature=1.0,
+                                                            top_k=None,
+                                                            top_p=None,
+                                                            no_repeat_ngram_size=3,
+                                                            num_return_sequences=1)
+            answer = eli5_tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,
+                                                 clean_up_tokenization_spaces=True)
+
+            kilt_example = create_kilt_datapoint(item["id"], query, answer[0], res_list)
+            kilt_output.append(kilt_example)
+            progress_bar.update(1)
+
+    with open(args.kilt_output_file, "w") as fp:
+        for kilt_example in kilt_output:
+            json.dump(kilt_example, fp)
+            fp.write("\n")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--kilt_input_file', default="./eli5-dev-kilt.jsonl", type=str)
+    parser.add_argument('--kilt_output_file', default="./eli5-predicted_retrieval.jsonl", type=str)
+    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",
+    )
+
+    args = parser.parse_args()
+
+    assert os.path.isfile(args.kilt_input_file), f"Input file {args.kilt_input_file} couldn't be loaded"
+    eval_generate(args)

util/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)

util/query_smoke_test.py

@@ -0,0 +1,81 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+from datasets import load_dataset
+
+
+def main():
+    device = ("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained('vblagoje/retribert-base-uncased')
+    model = AutoModel.from_pretrained('vblagoje/retribert-base-uncased').to(device)
+    _ = model.eval()
+
+    index_file_name = "./data/kilt_wikipedia.faiss"
+    kilt_wikipedia = load_dataset("kilt_wikipedia", split="full")
+    columns = ['kilt_id', 'wikipedia_id', 'wikipedia_title', 'text', 'anchors', 'categories',
+                              'wikidata_info', 'history']
+
+    min_snippet_length = 20
+    topk = 21
+
+    def articles_to_paragraphs(examples):
+        ids, titles, sections, texts, start_ps, end_ps, start_cs, end_cs = [], [], [], [], [], [], [], []
+        for bidx, example in enumerate(examples["text"]):
+            last_section = ""
+            for idx, p in enumerate(example["paragraph"]):
+                if "Section::::" in p:
+                    last_section = p
+                ids.append(examples["wikipedia_id"][bidx])
+                titles.append(examples["wikipedia_title"][bidx])
+                sections.append(last_section)
+                texts.append(p)
+                start_ps.append(idx)
+                end_ps.append(idx)
+                start_cs.append(0)
+                end_cs.append(len(p))
+
+        return {"wikipedia_id": ids, "title": titles,
+                "section": sections, "text": texts,
+                "start_paragraph_id": start_ps, "end_paragraph_id": end_ps,
+                "start_character": start_cs,
+                "end_character": end_cs
+                }
+
+    kilt_wikipedia_paragraphs = kilt_wikipedia.map(articles_to_paragraphs, batched=True,
+                                                   remove_columns=columns,
+                                                   batch_size=256, cache_file_name=f"./wiki_kilt_paragraphs_full.arrow",
+                                                   desc="Expanding wiki articles into paragraphs")
+
+    # use paragraphs that are not simple fragments or very short sentences
+    kilt_wikipedia_paragraphs = kilt_wikipedia_paragraphs.filter(lambda x: x["end_character"] > 250)
+    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", index_file_name, device=0)
+
+    def embed_questions_for_retrieval(questions):
+        query = tokenizer(questions, max_length=128, padding=True, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            q_reps = model.embed_questions(query["input_ids"].to(device),
+                                           query["attention_mask"].to(device)).cpu().type(torch.float)
+        return q_reps.numpy()
+
+    def query_index(question):
+        question_embedding = embed_questions_for_retrieval([question])
+        scores, wiki_passages = kilt_wikipedia_paragraphs.get_nearest_examples("embeddings", question_embedding, k=topk)
+        columns = ['wikipedia_id', 'title', 'text', 'section', 'start_paragraph_id', 'end_paragraph_id', 'start_character','end_character']
+        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
+
+    questions = ["What causes the contrails (cirrus aviaticus) behind jets at high altitude? ",
+                 "Why does water heated to a room temeperature feel colder than the air around it?"]
+    res_list = query_index(questions[0])
+    res_list = [res for res in res_list if len(res["text"].split()) > min_snippet_length][:int(topk / 3)]
+    for res in res_list:
+        print("\n")
+        print(res)
+
+
+main()
+
+