Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModel
3
4model = AutoModel.from_pretrained("ncbi/MedCPT-Article-Encoder")
5tokenizer = AutoTokenizer.from_pretrained("ncbi/MedCPT-Article-Encoder")
6
7# each article contains a list of two texts (usually a title and an abstract)
8articles = [
9 [
10 "Diagnosis and Management of Central Diabetes Insipidus in Adults",
11 "Central diabetes insipidus (CDI) is a clinical syndrome which results from loss or impaired function of vasopressinergic neurons in the hypothalamus/posterior pituitary, resulting in impaired synthesis and/or secretion of arginine vasopressin (AVP). [...]",
12 ],
13 [
14 "Adipsic diabetes insipidus",
15 "Adipsic diabetes insipidus (ADI) is a rare but devastating disorder of water balance with significant associated morbidity and mortality. Most patients develop the disease as a result of hypothalamic destruction from a variety of underlying etiologies. [...]",
16 ],
17 [
18 "Nephrogenic diabetes insipidus: a comprehensive overview",
19 "Nephrogenic diabetes insipidus (NDI) is characterized by the inability to concentrate urine that results in polyuria and polydipsia, despite having normal or elevated plasma concentrations of arginine vasopressin (AVP). [...]",
20 ],
21]
22
23with torch.no_grad():
24 # tokenize the articles
25 encoded = tokenizer(
26 articles,
27 truncation=True,
28 padding=True,
29 return_tensors='pt',
30 max_length=512,
31 )
32
33 # encode the queries (use the [CLS] last hidden states as the representations)
34 embeds = model(**encoded).last_hidden_state[:, 0, :]
35
36 print(embeds)
37 print(embeds.size())1tensor([[-0.0189, 0.0115, 0.0988, ..., -0.0655, 0.3155, -0.0357],
2 [-0.3402, -0.3064, -0.0749, ..., -0.0799, 0.3332, 0.1263],
3 [-0.2764, -0.0506, -0.0608, ..., 0.0389, 0.2532, 0.1580]])
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}