Function extract_bert

tools/bert/extract_bert_embeddings.py:33–48

(imdb_path, out_path, group_id=0, n_groups=1)

def extract_bert(imdb_path, out_path, group_id=0, n_groups=1):
    imdb = np.load(imdb_path)

    feat_extractor = BertFeatExtractor("bert-base-uncased")

    if group_id == 0:
        iterator_obj = tqdm(imdb[1:])
    else:
        iterator_obj = imdb[1:]

    for idx, el in enumerate(iterator_obj):
        if idx % n_groups != group_id:
            continue
        emb = feat_extractor.get_bert_embedding(el["question_str"])
        save_path = out_path + str(el["question_id"])
        np.save(save_path, emb.cpu().numpy())


if __name__ == "__main__":