Spaces:
Build error
Build error
File size: 4,884 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
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)
|