Spaces:
Build error
Build error
File size: 2,949 Bytes
e067d8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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)
|