forked from caodoanh2001/COSFormer
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathextract_word_embedding.py
More file actions
63 lines (53 loc) · 1.76 KB
/
extract_word_embedding.py
File metadata and controls
63 lines (53 loc) · 1.76 KB
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
from transformers import AutoTokenizer, AutoModel
import torch
import os
import json
import tqdm
# Mean Pooling - Take attention mask into account for correct averaging
def meanpooling(output, mask):
embeddings = output[0] # First element of model_output contains all token embeddings
mask = mask.unsqueeze(-1).expand(embeddings.size()).float()
return torch.sum(embeddings * mask, 1) / torch.clamp(mask.sum(1), min=1e-9)
# Sentences we want sentence embeddings for
word_dict = {
"[CLS]": 0,
"[UNK]": 1,
"[SEP]": 2,
"[PAD]": 3,
"non": 4,
"tumor.": 5,
"adenocarcinoma.": 6,
"squamous": 7,
"cell": 8,
"carcinoma.": 9,
"invasive": 10,
"ductal.": 11,
"lobular": 12,
"clear": 13,
"cell.": 14,
"papillary.": 15,
"chromophobe": 16,
"renal": 17,
"seminoma.": 18,
"mixed": 19,
"germ": 20,
}
task_name_embedding_path = './task_name_embeddings'
# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("neuml/pubmedbert-base-embeddings")
model = AutoModel.from_pretrained("neuml/pubmedbert-base-embeddings")
# tokenizer = AutoTokenizer.from_pretrained("dmis-lab/biobert-v1.1")
# model = AutoModel.from_pretrained("dmis-lab/biobert-v1.1")
list_embeddings = []
for i, task in tqdm.tqdm(enumerate(word_dict)):
text = task
# Tokenize sentences
inputs = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
# Compute token embeddings
with torch.no_grad():
output = model(**inputs)
embeddings = output['pooler_output']
task_name_embedding_path_name = os.path.join(task_name_embedding_path, task + '.pth')
list_embeddings.append(embeddings)
word_embed = torch.stack(list_embeddings).squeeze(1)
torch.save(word_embed, "TCGA_word_embed.pth")