|
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") |
|
|
|
|
|
|
|
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) |
|
|
