Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4tokenizer = AutoTokenizer.from_pretrained("ncbi/MedCPT-Cross-Encoder")
5model = AutoModelForSequenceClassification.from_pretrained("ncbi/MedCPT-Cross-Encoder")
6
7query = "diabetes treatment"
8
9# 6 articles to be ranked for the input query
10articles = [
11 "Type 1 and 2 diabetes mellitus: A review on current treatment approach and gene therapy as potential intervention. Type 1 and type 2 diabetes mellitus is a serious and lifelong condition commonly characterised by abnormally elevated blood glucose levels due to a failure in insulin production or a decrease in insulin sensitivity and function. [...]",
12 "Diabetes mellitus and its chronic complications. Diabetes mellitus is a major cause of morbidity and mortality, and it is a major risk factor for early onset of coronary heart disease. Complications of diabetes are retinopathy, nephropathy, and peripheral neuropathy. [...]",
13 "Diagnosis and Management of Central Diabetes Insipidus in Adults. 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). [...]",
14 "Adipsic diabetes insipidus. 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. [...]",
15 "Nephrogenic diabetes insipidus: a comprehensive overview. 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). [...]",
16 "Impact of Salt Intake on the Pathogenesis and Treatment of Hypertension. Excessive dietary salt (sodium chloride) intake is associated with an increased risk for hypertension, which in turn is especially a major risk factor for stroke and other cardiovascular pathologies, but also kidney diseases. Besides, high salt intake or preference for salty food is discussed to be positive associated with stomach cancer, and according to recent studies probably also obesity risk. [...]"
17]
18
19# combine query article into pairs
20pairs = [[query, article] for article in articles]
21
22with torch.no_grad():
23 encoded = tokenizer(
24 pairs,
25 truncation=True,
26 padding=True,
27 return_tensors="pt",
28 max_length=512,
29 )
30
31 logits = model(**encoded).logits.squeeze(dim=1)
32
33 print(logits)tensor([ 6.9363, -8.2063, -8.7692, -12.3450, -10.4416, -15.8475])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}