Upload create_faiss_index.py

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)