Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModel
3
4model = AutoModel.from_pretrained("ncbi/MedCPT-Query-Encoder")
5tokenizer = AutoTokenizer.from_pretrained("ncbi/MedCPT-Query-Encoder")
6
7queries = [
8 "diabetes treatment",
9 "How to treat diabetes?",
10 "A 45-year-old man presents with increased thirst and frequent urination over the past 3 months.",
11]
12
13with torch.no_grad():
14 # tokenize the queries
15 encoded = tokenizer(
16 queries,
17 truncation=True,
18 padding=True,
19 return_tensors='pt',
20 max_length=64,
21 )
22
23 # encode the queries (use the [CLS] last hidden states as the representations)
24 embeds = model(**encoded).last_hidden_state[:, 0, :]
25
26 print(embeds)
27 print(embeds.size())1tensor([[ 0.0413, 0.0084, -0.0491, ..., -0.4963, -0.3830, -0.3593],
2 [ 0.0801, 0.1193, -0.0905, ..., -0.5380, -0.5059, -0.2944],
3 [-0.3412, 0.1521, -0.0946, ..., 0.0952, 0.1660, -0.0902]])
4torch.Size([3, 768])1@article{jin2023medcpt,
2 title={MedCPT: Contrastive Pre-trained Transformers with large-scale PubMed search logs for zero-shot biomedical information retrieval},
3 author={Jin, Qiao and Kim, Won and Chen, Qingyu and Comeau, Donald C and Yeganova, Lana and Wilbur, W John and Lu, Zhiyong},
4 journal={Bioinformatics},
5 volume={39},
6 number={11},
7 pages={btad651},
8 year={2023},
9 publisher={Oxford University Press}
10}