def do_combine(): """Combines all shards of the fact2fact matrix.""" lines = _load_lines() num_facts = len(lines) rows, cols, vals = [], [], [] shard_range = list(range(FLAGS.num_shards + 1)) for shard_id in tqdm(shard_range, desc="loading shard files"): tf.logging.info("Appending %d-th shard ...", shard_id) with tf.gfile.Open( os.path.join(FLAGS.fact2fact_index_dir, "f2f_%d.json" % shard_id)) as f: cur_data = json.load(f) rows += cur_data["rows"] cols += cur_data["cols"] vals += cur_data["vals"] del cur_data gc.collect() tf.logging.info("Creating sp.csr_matrix") sp_fact2fact = sp.csr_matrix((vals, (rows, cols)), shape=[num_facts, num_facts]) tf.logging.info("Num nonzero in f2f = %d", sp_fact2fact.getnnz()) tf.logging.info("Saving as ragged f2f tensor %s.", str(sp_fact2fact.shape)) search_utils.write_ragged_to_checkpoint( "fact2fact", sp_fact2fact, os.path.join(FLAGS.fact2fact_index_dir, "fact2fact.npz"))
def main(_): if not tf.gfile.Exists(FLAGS.multihop_output_dir): tf.gfile.MakeDirs(FLAGS.multihop_output_dir) # Filenames. paragraphs_file = os.path.join(FLAGS.data_dir, "processed_wiki.json") train_file = os.path.join(FLAGS.qry_dir, "train.json") dev_file = os.path.join(FLAGS.qry_dir, "dev.json") test_file = os.path.join(FLAGS.qry_dir, "test.json") entities_file = os.path.join(FLAGS.data_dir, "entities.txt") # Initialize tokenizer. tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case) # Read entities. tf.logging.info("Reading entities.") entity2id, entity2name = {}, {} with tf.gfile.Open(entities_file) as f: for ii, line in tqdm(enumerate(f)): entity = line.strip() entity2id[entity] = ii entity2name[entity] = entity # Read paragraphs, mentions and entities. mentions = [] ent_rows, ent_cols, ent_vals = [], [], [] para_rows, para_cols, para_vals = [], [], [] mention2text = {} total_sub_paras = [0] all_sub_paras = [] entity_counts = collections.defaultdict(int) tf.logging.info("Reading paragraphs from %s", paragraphs_file) with tf.gfile.Open(paragraphs_file) as f: for line in tqdm(f): orig_para = json.loads(line.strip()) sub_para_objs = _get_sub_paras(orig_para, tokenizer, FLAGS.max_seq_length, FLAGS.doc_stride, total_sub_paras) for para_obj in sub_para_objs: # Add mentions from this paragraph. my_entities = [] my_mentions = [] for m in para_obj["mentions"]: # Para to mention matrix. para_rows.append(para_obj["id"]) para_cols.append(len(mentions)) para_vals.append(1.) # Create entity to mention sparse connections. my_entities.append(m["kb_id"]) my_mentions.append(len(mentions)) entity_counts[m["kb_id"]] += 1 mention2text[len(mentions)] = m["text"] mentions.append((m["kb_id"], para_obj["id"], m["start_token"], m["end_token"])) for entity in my_entities: ent_rows.append(entity) ent_cols.append(para_obj["id"]) ent_vals.append(1. / len(my_mentions)) all_sub_paras.append(para_obj["tokens"]) assert len(all_sub_paras) == total_sub_paras[0], ( len(all_sub_paras), total_sub_paras) tf.logging.info("Num paragraphs = %d, Num mentions = %d", total_sub_paras[0], len(mentions)) tf.logging.info("Saving coreference map.") search_utils.write_to_checkpoint( "coref", np.array([m[0] for m in mentions], dtype=np.int32), tf.int32, os.path.join(FLAGS.multihop_output_dir, "coref.npz")) tf.logging.info("Creating entity to mentions matrix.") sp_entity2para = sp.csr_matrix((ent_vals, (ent_rows, ent_cols)), shape=[len(entity2id), len(all_sub_paras)]) sp_entity2para_filt = preprocess_utils.filter_sparse_rows( sp_entity2para, FLAGS.max_paragraphs_per_entity) sp_para2ment = sp.csr_matrix((para_vals, (para_rows, para_cols)), shape=[len(all_sub_paras), len(mentions)]) sp_entity2mention = sp_entity2para_filt.dot(sp_para2ment) tf.logging.info("Num nonzero = %d", sp_entity2mention.getnnz()) tf.logging.info("Saving as ragged tensor %s.", str(sp_entity2mention.shape)) search_utils.write_ragged_to_checkpoint( "ent2ment", sp_entity2mention, os.path.join(FLAGS.multihop_output_dir, "ent2ment.npz")) tf.logging.info("Saving mentions metadata.") np.save( tf.gfile.Open(os.path.join(FLAGS.multihop_output_dir, "mentions.npy"), "w"), np.array(mentions, dtype=np.int64)) json.dump( mention2text, tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "mention2text.json"), "w")) tf.logging.info("Saving entities metadata.") json.dump([entity2id, entity2name], tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "entities.json"), "w")) json.dump( entity_counts, tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "entity_counts.json"), "w")) tf.logging.info("Saving split paragraphs.") json.dump( all_sub_paras, tf.gfile.Open(os.path.join(FLAGS.multihop_output_dir, "subparas.json"), "w")) # Store entity tokens. tf.logging.info("Processing entities.") entity_ids = np.zeros((len(entity2id), FLAGS.max_entity_len), dtype=np.int32) entity_mask = np.zeros((len(entity2id), FLAGS.max_entity_len), dtype=np.float32) num_exceed_len = 0. for entity in tqdm(entity2id): ei = entity2id[entity] entity_tokens = tokenizer.tokenize(entity2name[entity]) entity_token_ids = tokenizer.convert_tokens_to_ids(entity_tokens) if len(entity_token_ids) > FLAGS.max_entity_len: num_exceed_len += 1 entity_token_ids = entity_token_ids[:FLAGS.max_entity_len] entity_ids[ei, :len(entity_token_ids)] = entity_token_ids entity_mask[ei, :len(entity_token_ids)] = 1. tf.logging.info( "Saving %d entity ids and mask. %d exceed max-length of %d.", len(entity2id), num_exceed_len, FLAGS.max_entity_len) search_utils.write_to_checkpoint( "entity_ids", entity_ids, tf.int32, os.path.join(FLAGS.multihop_output_dir, "entity_ids")) search_utils.write_to_checkpoint( "entity_mask", entity_mask, tf.float32, os.path.join(FLAGS.multihop_output_dir, "entity_mask")) # Pre-process question files. def _preprocess_qrys(in_file, out_file): tf.logging.info("Working on %s", in_file) with tf.gfile.Open(in_file) as f_in, tf.gfile.Open(out_file, "w") as f_out: for line in f_in: item = json.loads(line.strip()) # Sort entities in ascending order of their frequencies. e_counts = [ entity_counts[e["kb_id"]] for e in item["entities"] ] sorted_i = sorted(enumerate(e_counts), key=lambda x: x[1]) item["entities"] = [item["entities"][ii] for ii, _ in sorted_i] f_out.write(json.dumps(item) + "\n") _preprocess_qrys(train_file, os.path.join(FLAGS.multihop_output_dir, "train.json")) _preprocess_qrys(dev_file, os.path.join(FLAGS.multihop_output_dir, "dev.json")) _preprocess_qrys(test_file, os.path.join(FLAGS.multihop_output_dir, "test.json")) # Copy BERT checkpoint for future use. tf.logging.info("Copying BERT checkpoint.") if tf.gfile.Exists(os.path.join(FLAGS.pretrain_dir, "best_model.index")): bert_ckpt = os.path.join(FLAGS.pretrain_dir, "best_model") else: bert_ckpt = tf.train.latest_checkpoint(FLAGS.pretrain_dir) tf.logging.info("%s.data-00000-of-00001", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".data-00000-of-00001", os.path.join(FLAGS.multihop_output_dir, "bert_init.data-00000-of-00001"), overwrite=True) tf.logging.info("%s.index", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".index", os.path.join(FLAGS.multihop_output_dir, "bert_init.index"), overwrite=True) tf.logging.info("%s.meta", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".meta", os.path.join(FLAGS.multihop_output_dir, "bert_init.meta"), overwrite=True) # Get mention embeddings from BERT. tf.logging.info("Computing mention embeddings for %d paras.", len(all_sub_paras)) bert_predictor = bert_utils.BERTPredictor(tokenizer, bert_ckpt) para_emb = bert_predictor.get_doc_embeddings(all_sub_paras) mention_emb = np.empty((len(mentions), 2 * bert_predictor.emb_dim), dtype=np.float32) for im, mention in enumerate(mentions): mention_emb[im, :] = np.concatenate([ para_emb[mention[1], mention[2], :], para_emb[mention[1], mention[3], :] ]) del para_emb tf.logging.info("Saving %d mention features to tensorflow checkpoint.", mention_emb.shape[0]) with tf.device("/cpu:0"): search_utils.write_to_checkpoint( "db_emb", mention_emb, tf.float32, os.path.join(FLAGS.multihop_output_dir, "mention_feats"))
def do_preprocess(tokenizer): """Loads and processes the data.""" # Read concepts. tf.logging.info("Reading entities.") entity2id = load_concept_vocab(FLAGS.entity_file) entity2name = {concept: concept for (concept, _) in entity2id.items()} # print("print # concepts:", len(entity2id)) tf.logging.info("# concepts: %d", len(entity2id)) if not tf.gfile.Exists(FLAGS.multihop_output_dir): tf.gfile.MakeDirs(FLAGS.multihop_output_dir) # Read paragraphs, mentions and entities. mentions = [] ent_rows, ent_cols, ent_vals = [], [], [] ent2fact_rows, ent2fact_cols, ent2fact_vals = [], [], [] fact2ent_rows, fact2ent_cols, fact2ent_vals = [], [], [] ent2num_facts = collections.defaultdict(lambda: 0) mention2text = {} total_sub_paras = [0] all_sub_paras = [] num_skipped_mentions = 0. tf.logging.info("Reading paragraphs from %s", FLAGS.wiki_file) with tf.gfile.Open(FLAGS.wiki_file) as f: lines = f.read().split("\n") if not lines[-1]: lines = lines[:-1] fact2entity = [] for ii, line in tqdm(enumerate(lines[:]), total=len(lines), desc="preprocessing lines"): if ii == FLAGS.max_total_paragraphs: tf.logging.info( "Processed maximum number of paragraphs, breaking.") break orig_para = json.loads(line.strip()) if orig_para["kb_id"].lower() not in entity2id: tf.logging.info("%s not in entities. Skipping %s para", orig_para["kb_id"], orig_para["title"]) continue sub_para_objs = index_util.get_sub_paras(orig_para, tokenizer, FLAGS.max_seq_length, FLAGS.doc_stride, total_sub_paras) assert len(sub_para_objs) == 1 # each doc is a single paragraph. for para_obj in sub_para_objs: # Add mentions from this paragraph. local2global = {} # title_entity_mention = None assert para_obj["id"] == len(fact2entity) fact2entity.append([]) for im, mention in enumerate( para_obj["mentions"][:FLAGS.max_mentions_per_doc]): if mention["kb_id"].lower() not in entity2id: # tf.logging.info("%s not in entities. Skipping mention %s", # mention["kb_id"], mention["text"]) num_skipped_mentions += 1 continue mention2text[len(mentions)] = mention["text"] # Map the index of a local mention to the global mention id. local2global[im] = len(mentions) # if mention["kb_id"] == orig_para["kb_id"]: # title_entity_mention = len(mentions) # The shape of 'mentions.npy' is thus #mentions * 4. mentions.append( (entity2id[mention["kb_id"].lower()], para_obj["id"], mention["start_token"], mention["end_token"])) # fact_id 2 entity_ids fact2entity[para_obj["id"]].append( entity2id[mention["kb_id"].lower()]) # Note: each pair of mention in this paragraph is recorded. local_mentioned_entity_ids = set( [mentions[gm][0] for _, gm in local2global.items()]) # Creating sparse entries for entity2mention matrix. for _, gm in local2global.items(): for cur_entity_id in local_mentioned_entity_ids: if cur_entity_id != gm: ent_rows.append(cur_entity_id) ent_cols.append(gm) ent_vals.append(1.) # Creating sparse entries for entity2fact matrix. for cur_entity_id in local_mentioned_entity_ids: fact2ent_rows.append(ii) # doc_id fact2ent_cols.append(cur_entity_id) fact2ent_vals.append(1.) # Note: Use tf-idf to limit this in the future. if ent2num_facts[cur_entity_id] >= FLAGS.max_facts_per_entity: # We want to limit the number of the facts in ent2fact. # Otherwise, the init_fact will be huge. continue ent2fact_rows.append(cur_entity_id) ent2fact_cols.append(ii) # doc_id ent2fact_vals.append(1.) ent2num_facts[cur_entity_id] += 1 all_sub_paras.append(para_obj["tokens"]) assert len(all_sub_paras) == total_sub_paras[0], (len(all_sub_paras), total_sub_paras) tf.logging.info("Num paragraphs = %d, Num mentions = %d", total_sub_paras[0], len(mentions)) tf.logging.info("Saving mention2entity coreference map.") search_utils.write_to_checkpoint( "coref", np.array([m[0] for m in mentions], dtype=np.int32), tf.int32, os.path.join(FLAGS.multihop_output_dir, "coref.npz")) tf.logging.info("Creating ent2men matrix with %d entries.", len(ent_vals)) # Fill a zero-inited sparse matrix for entity2mention. sp_entity2mention = sp.csr_matrix((ent_vals, (ent_rows, ent_cols)), shape=[len(entity2id), len(mentions)]) tf.logging.info("Num nonzero in e2m = %d", sp_entity2mention.getnnz()) tf.logging.info("Saving as ragged e2m tensor %s.", str(sp_entity2mention.shape)) search_utils.write_ragged_to_checkpoint( "ent2ment", sp_entity2mention, os.path.join(FLAGS.multihop_output_dir, "ent2ment.npz")) tf.logging.info("Saving mentions metadata.") np.save( tf.gfile.Open(os.path.join(FLAGS.multihop_output_dir, "mentions.npy"), "w"), np.array(mentions, dtype=np.int64)) json.dump( mention2text, tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "mention2text.json"), "w")) tf.logging.info("Saving entities metadata.") assert len(lines) == len(all_sub_paras) num_facts = len(all_sub_paras) # Fill a zero-inited sparse matrix for entity2fact. sp_entity2fact = sp.csr_matrix( (ent2fact_vals, (ent2fact_rows, ent2fact_cols)), shape=[len(entity2id), num_facts]) tf.logging.info("Num nonzero in e2f = %d", sp_entity2fact.getnnz()) tf.logging.info("Saving as ragged e2f tensor %s.", str(sp_entity2fact.shape)) search_utils.write_ragged_to_checkpoint( "ent2fact", sp_entity2fact, os.path.join(FLAGS.multihop_output_dir, "ent2fact_%d.npz" % FLAGS.max_facts_per_entity)) # Fill a zero-inited sparse matrix for fact2entity. sp_fact2entity = sp.csr_matrix( (ent2fact_vals, (ent2fact_cols, ent2fact_rows)), # Transpose. shape=[num_facts, len(entity2id)]) tf.logging.info("Num nonzero in f2e = %d", sp_fact2entity.getnnz()) tf.logging.info("Saving as ragged f2e tensor %s.", str(sp_fact2entity.shape)) search_utils.write_ragged_to_checkpoint( "fact2ent", sp_fact2entity, os.path.join(FLAGS.multihop_output_dir, "fact_coref.npz")) json.dump([entity2id, entity2name], tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "entities.json"), "w")) tf.logging.info("Saving split paragraphs.") json.dump( all_sub_paras, tf.gfile.Open(os.path.join(FLAGS.multihop_output_dir, "subparas.json"), "w")) # Store entity tokens. tf.logging.info("Processing entities.") entity_ids = np.zeros((len(entity2id), FLAGS.max_entity_length), dtype=np.int32) entity_mask = np.zeros((len(entity2id), FLAGS.max_entity_length), dtype=np.float32) num_exceed_len = 0. for entity in tqdm(entity2id): ei = entity2id[entity] entity_tokens = tokenizer.tokenize(entity2name[entity]) entity_token_ids = tokenizer.convert_tokens_to_ids(entity_tokens) if len(entity_token_ids) > FLAGS.max_entity_length: num_exceed_len += 1 entity_token_ids = entity_token_ids[:FLAGS.max_entity_length] entity_ids[ei, :len(entity_token_ids)] = entity_token_ids entity_mask[ei, :len(entity_token_ids)] = 1. tf.logging.info("Saving %d entity ids. %d exceed max-length of %d.", len(entity2id), num_exceed_len, FLAGS.max_entity_length) search_utils.write_to_checkpoint( "entity_ids", entity_ids, tf.int32, os.path.join(FLAGS.multihop_output_dir, "entity_ids")) search_utils.write_to_checkpoint( "entity_mask", entity_mask, tf.float32, os.path.join(FLAGS.multihop_output_dir, "entity_mask"))
def main(_): if not tf.gfile.Exists(FLAGS.multihop_output_dir): tf.gfile.MakeDirs(FLAGS.multihop_output_dir) # Initialize tokenizer. tokenizer = tokenization.FullTokenizer(vocab_file=FLAGS.vocab_file, do_lower_case=FLAGS.do_lower_case) # Read entities. if FLAGS.do_preprocess: tf.logging.info("Reading entities.") entity2id, entity2name = {}, {} with tf.gfile.Open(FLAGS.entity_file) as f: entities = json.load(f) tf.logging.info("Read %d entities", len(entities)) for e, (_, n) in entities.items(): if e.lower() in entity2id: continue # tf.logging.warn("%s entity repeated", e) entity2id[e.lower()] = len(entity2id) entity2name[e.lower()] = n tf.logging.info("Kept %d entities", len(entity2id)) # Read paragraphs, mentions and entities. if FLAGS.do_preprocess: mentions = [] ent_rows, ent_cols, ent_vals = [], [], [] mention2text = {} total_sub_paras = [0] all_sub_paras = [] num_skipped_mentions = 0. tf.logging.info("Reading paragraphs from %s", FLAGS.wiki_file) with tf.gfile.Open(FLAGS.wiki_file) as f: for ii, line in tqdm(enumerate(f)): if ii == FLAGS.max_total_paragraphs: tf.logging.info( "Processed maximum number of paragraphs, breaking.") break if ii > 0 and ii % 100000 == 0: tf.logging.info("Skipped / Kept mentions = %.3f", num_skipped_mentions / len(mentions)) orig_para = json.loads(line.strip()) if orig_para["kb_id"].lower() not in entity2id: tf.logging.warn("%s not in entities. Skipping %s para", orig_para["kb_id"], orig_para["title"]) continue sub_para_objs = _get_sub_paras(orig_para, tokenizer, FLAGS.max_seq_length, FLAGS.doc_stride, total_sub_paras) for para_obj in sub_para_objs: # Add mentions from this paragraph. local2global = {} title_entity_mention = None for im, mention in enumerate( para_obj["mentions"] [:FLAGS.max_mentions_per_entity]): if mention["kb_id"].lower() not in entity2id: # tf.logging.warn("%s not in entities. Skipping mention %s", # mention["kb_id"], mention["text"]) num_skipped_mentions += 1 continue mention2text[len(mentions)] = mention["text"] local2global[im] = len(mentions) if mention["kb_id"] == orig_para["kb_id"]: title_entity_mention = len(mentions) mentions.append( (entity2id[mention["kb_id"].lower()], para_obj["id"], mention["start_token"], mention["end_token"])) for im, gm in local2global.items(): # entity to mention matrix. ent_rows.append(entity2id[orig_para["kb_id"].lower()]) ent_cols.append(gm) ent_vals.append(1.) if title_entity_mention is not None: ent_rows.append(mentions[gm][0]) ent_cols.append(title_entity_mention) ent_vals.append(1.) all_sub_paras.append(para_obj["tokens"]) assert len(all_sub_paras) == total_sub_paras[0], ( len(all_sub_paras), total_sub_paras) tf.logging.info("Num paragraphs = %d, Num mentions = %d", total_sub_paras[0], len(mentions)) tf.logging.info("Saving coreference map.") search_utils.write_to_checkpoint( "coref", np.array([m[0] for m in mentions], dtype=np.int32), tf.int32, os.path.join(FLAGS.multihop_output_dir, "coref.npz")) tf.logging.info("Creating entity to mentions matrix.") sp_entity2mention = sp.csr_matrix( (ent_vals, (ent_rows, ent_cols)), shape=[len(entity2id), len(mentions)]) tf.logging.info("Num nonzero = %d", sp_entity2mention.getnnz()) tf.logging.info("Saving as ragged tensor %s.", str(sp_entity2mention.shape)) search_utils.write_ragged_to_checkpoint( "ent2ment", sp_entity2mention, os.path.join(FLAGS.multihop_output_dir, "ent2ment.npz")) tf.logging.info("Saving mentions metadata.") np.save( tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "mentions.npy"), "w"), np.array(mentions, dtype=np.int64)) json.dump( mention2text, tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "mention2text.json"), "w")) tf.logging.info("Saving entities metadata.") json.dump([entity2id, entity2name], tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "entities.json"), "w")) tf.logging.info("Saving split paragraphs.") json.dump( all_sub_paras, tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "subparas.json"), "w")) # Store entity tokens. if FLAGS.do_preprocess: tf.logging.info("Processing entities.") entity_ids = np.zeros((len(entity2id), FLAGS.max_entity_length), dtype=np.int32) entity_mask = np.zeros((len(entity2id), FLAGS.max_entity_length), dtype=np.float32) num_exceed_len = 0. for entity in tqdm(entity2id): ei = entity2id[entity] entity_tokens = tokenizer.tokenize(entity2name[entity]) entity_token_ids = tokenizer.convert_tokens_to_ids(entity_tokens) if len(entity_token_ids) > FLAGS.max_entity_length: num_exceed_len += 1 entity_token_ids = entity_token_ids[:FLAGS.max_entity_length] entity_ids[ei, :len(entity_token_ids)] = entity_token_ids entity_mask[ei, :len(entity_token_ids)] = 1. tf.logging.info("Saving %d entity ids. %d exceed max-length of %d.", len(entity2id), num_exceed_len, FLAGS.max_entity_length) search_utils.write_to_checkpoint( "entity_ids", entity_ids, tf.int32, os.path.join(FLAGS.multihop_output_dir, "entity_ids")) search_utils.write_to_checkpoint( "entity_mask", entity_mask, tf.float32, os.path.join(FLAGS.multihop_output_dir, "entity_mask")) # Copy BERT checkpoint for future use. if FLAGS.do_preprocess: tf.logging.info("Copying BERT checkpoint.") if tf.gfile.Exists(os.path.join(FLAGS.pretrain_dir, "best_model.index")): bert_ckpt = os.path.join(FLAGS.pretrain_dir, "best_model") else: bert_ckpt = tf.train.latest_checkpoint(FLAGS.pretrain_dir) tf.logging.info("%s.data-00000-of-00001", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".data-00000-of-00001", os.path.join(FLAGS.multihop_output_dir, "bert_init.data-00000-of-00001"), overwrite=True) tf.logging.info("%s.index", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".index", os.path.join(FLAGS.multihop_output_dir, "bert_init.index"), overwrite=True) tf.logging.info("%s.meta", bert_ckpt) tf.gfile.Copy(bert_ckpt + ".meta", os.path.join(FLAGS.multihop_output_dir, "bert_init.meta"), overwrite=True) if FLAGS.do_embed: # Get mention embeddings from BERT. bert_ckpt = os.path.join(FLAGS.multihop_output_dir, "bert_init") if not FLAGS.do_preprocess: with tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "mentions.npy")) as f: mentions = np.load(f) with tf.gfile.Open( os.path.join(FLAGS.multihop_output_dir, "subparas.json")) as f: all_sub_paras = json.load(f) tf.logging.info("Computing embeddings for %d mentions over %d paras.", len(mentions), len(all_sub_paras)) shard_size = len(mentions) // FLAGS.num_shards bert_predictor = bert_utils_v2.BERTPredictor(tokenizer, bert_ckpt) if FLAGS.my_shard is None: shard_range = range(FLAGS.num_shards + 1) else: shard_range = [FLAGS.my_shard] for ns in shard_range: min_ = ns * shard_size max_ = (ns + 1) * shard_size if min_ >= len(mentions): break if max_ > len(mentions): max_ = len(mentions) min_subp = mentions[min_][1] max_subp = mentions[max_ - 1][1] tf.logging.info("Processing shard %d of %d mentions and %d paras.", ns, max_ - min_, max_subp - min_subp + 1) para_emb = bert_predictor.get_doc_embeddings( all_sub_paras[min_subp:max_subp + 1]) assert para_emb.shape[2] == 2 * FLAGS.projection_dim mention_emb = np.empty((max_ - min_, 2 * bert_predictor.emb_dim), dtype=np.float32) for im, mention in enumerate(mentions[min_:max_]): mention_emb[im, :] = np.concatenate([ para_emb[mention[1] - min_subp, mention[2], :FLAGS.projection_dim], para_emb[mention[1] - min_subp, mention[3], FLAGS.projection_dim:2 * FLAGS.projection_dim] ]) del para_emb tf.logging.info( "Saving %d mention features to tensorflow checkpoint.", mention_emb.shape[0]) with tf.device("/cpu:0"): search_utils.write_to_checkpoint( "db_emb_%d" % ns, mention_emb, tf.float32, os.path.join(FLAGS.multihop_output_dir, "mention_feats_%d" % ns)) if FLAGS.do_combine: # Combine sharded DB into one. if FLAGS.shards_to_combine is None: shard_range = range(FLAGS.num_shards + 1) else: shard_range = range(FLAGS.shards_to_combine) with tf.device("/cpu:0"): all_db = [] for i in shard_range: ckpt_path = os.path.join(FLAGS.multihop_output_dir, "mention_feats_%d" % i) reader = tf.NewCheckpointReader(ckpt_path) var_to_shape_map = reader.get_variable_to_shape_map() tf.logging.info("Reading %s from %s with shape %s", "db_emb_%d" % i, ckpt_path, str(var_to_shape_map["db_emb_%d" % i])) tf_db = search_utils.load_database( "db_emb_%d" % i, var_to_shape_map["db_emb_%d" % i], ckpt_path) all_db.append(tf_db) tf.logging.info("Reading all variables.") session = tf.Session() session.run(tf.global_variables_initializer()) session.run(tf.local_variables_initializer()) np_db = session.run(all_db) tf.logging.info("Concatenating and storing.") np_db = np.concatenate(np_db, axis=0) search_utils.write_to_checkpoint( "db_emb", np_db, tf.float32, os.path.join(FLAGS.multihop_output_dir, "mention_feats"))