-
Notifications
You must be signed in to change notification settings - Fork 3
/
generate_embeddings.py
26 lines (20 loc) · 1.04 KB
/
generate_embeddings.py
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
import argparse
import numpy as np
from cord_ann.embeddings import EmbeddingModel
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--sentences_path', default="cord_19_dataset_formatted/cord_19_sentences.txt",
help='Path to extracted sentences')
parser.add_argument('--embedding_path', default="embeddings.npy",
help='Output path of the generated embeddings')
parser.add_argument('--model_name_or_path', default='bert-base-nli-mean-tokens')
parser.add_argument('--batch_size', default=8, type=int)
parser.add_argument('--device', default='cuda')
args = parser.parse_args()
with open(args.sentences_path) as f:
sentences = f.read().split('\n')
model = EmbeddingModel(model_name_or_path=args.model_name_or_path,
device=args.device,
batch_size=args.batch_size)
sentence_embeddings = model.encode_sentences(sentences=sentences)
np.save(args.embedding_path, sentence_embeddings)