Views
No views yet

1from datasets import load_dataset
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4
5# Load base model
6base_model = AutoModelForCausalLM.from_pretrained(
7 'mistralai/Mistral-7B-v0.1',
8 trust_remote_code=True,
9 device_map="auto",
10 torch_dtype=torch.bfloat16, # Optional, if you have insufficient VRAM, lower the precision.
11)
12
13# Load tokenizer
14tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-7B-v0.1')
15tokenizer.add_special_tokens(dict(
16 eos_token=AddedToken("<|im_end|>", single_word=False, lstrip=False, rstrip=False, normalized=True, special=True),
17 unk_token=AddedToken("<unk>", single_word=False, lstrip=False, rstrip=False, normalized=True, special=True),
18 bos_token=AddedToken("<s>", single_word=False, lstrip=False, rstrip=False, normalized=True, special=True),
19 pad_token=AddedToken("</s>", single_word=False, lstrip=False, rstrip=False, normalized=False, special=True),
20))
21tokenizer.add_tokens([AddedToken("<|im_start|>", single_word=False, lstrip=True, rstrip=True, normalized=False)], special_tokens=True)
22tokenizer.additional_special_tokens = ['<unk>', '<s>', '</s>', '<|im_end|>', '<|im_start|>']
23
24model.resize_token_embeddings(len(tokenizer))
25model.config.eos_token_id = tokenizer.eos_token_id
26
27# Load PEFT
28model = PeftModel.from_pretrained(base_model, 'Lowenzahn/PathoIE-Mistral-7B')
29model = model.merge_and_unload()
30model = model.eval()
31
32# Inference
33prompts = ["Machine learning is"]
34inputs = tokenizer(prompts, return_tensors="pt")
35gen_kwargs = {"max_new_tokens": 1024, "top_p": 0.8, "temperature": 0.0, "do_sample": False, "repetition_penalty": 1.0}
36output = model.generate(inputs['input_ids'], **gen_kwargs)
37output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
38print(output)<|im_start|> system
You are a pathologist who specialized in lung cancer.
Your task is extracting informations requested by the user from the lung cancer pathology report and formatting extracted informations into JSON.
The information to be extracted is clearly specified in the report, so one must avoid from inferring information that is not present.
Remember, you MUST answer in JSON only. Avoid any additional explanations.
<|im_end|>
<|im_start|> user
Extract the following informations (value-set) from the report I provide.
If the required information to extract each value in the value-set is not present in the pathology report, consider it as 'not submitted'.
<value-set>
- MORPHOLOGY_DIAGNOSIS
- SUBTYPE_DOMINANT
- MAX_SIZE_OF_TUMOR(invasive component only)
- MAX_SIZE_OF_TUMOR(including CIS=AIS)
- INVASION_TO_VISCERAL_PLEURAL
- MAIN_BRONCHUS
- INVASION_TO_CHEST_WALL
- INVASION_TO_PARIETAL_PLEURA
- INVASION_TO_PERICARDIUM
- INVASION_TO_PHRENIC_NERVE
- TUMOR_SIZE_CNT
- LUNG_TO_LUNG_METASTASIS
- INTRAPULMONARY_METASTASIS
- SATELLITE_TUMOR_LOCATION
- SEPARATE_TUMOR_LOCATION
- INVASION_TO_MEDIASTINUM
- INVASION_TO_DIAPHRAGM
- INVASION_TO_HEART
- INVASION_TO_RECURRENT_LARYNGEAL_NERVE
- INVASION_TO_TRACHEA
- INVASION_TO_ESOPHAGUS
- INVASION_TO_SPINE
- METASTATIC_RIGHT_UPPER_LOBE
- METASTATIC_RIGHT_MIDDLE_LOBE
- METASTATIC_RIGHT_LOWER_LOBE
- METASTATIC_LEFT_UPPER_LOBE
- METASTATIC_LEFT_LOWER_LOBE
- INVASION_TO_AORTA
- INVASION_TO_SVC
- INVASION_TO_IVC
- INVASION_TO_PULMONARY_ARTERY
- INVASION_TO_PULMONARY_VEIN
- INVASION_TO_CARINA
- PRIMARY_CANCER_LOCATION_RIGHT_UPPER_LOBE
- PRIMARY_CANCER_LOCATION_RIGHT_MIDDLE_LOBE
- PRIMARY_CANCER_LOCATION_RIGHT_LOWER_LOBE
- PRIMARY_CANCER_LOCATION_LEFT_UPPER_LOBE
- PRIMARY_CANCER_LOCATION_LEFT_LOWER_LOBE
- RELATED_TO_ATELECTASIS_OR_OBSTRUCTIVE_PNEUMONITIS
- PRIMARY_SITE_LATERALITY
- LYMPH_METASTASIS_SITES
- NUMER_OF_LYMPH_NODE_META_CASES
---
<report>
[A] Lung, left lower lobe, lobectomy
1. ADENOSQUAMOUS CARCINOMA [by 2015 WHO classification]
- other subtype: acinar (50%), lepidic (30%), solid (20%)
1) Pre-operative / Previous treatment: not done
2) Histologic grade: moderately differentiated
3) Size of tumor:
a. Invasive component only: 3.5 x 2.5 x 1.3 cm, 2.4 x 2.3 x 1.1 cm
b. Including CIS component: 3.9 x 2.6 x 1.3 cm, 3.8 x 3.1 x 1.2 cm
4) Extent of invasion
a. Invasion to visceral pleura: PRESENT (P2)
b. Invasion to superior vena cava: present
5) Main bronchus: not submitted
6) Necrosis: absent
7) Resection margin: free from carcinoma (safey margin: 1.1 cm)
8) Lymph node: metastasis in 2 out of 10 regional lymph nodes
(peribronchial lymph node: 1/3, LN#5,6 :0/1, LN#7:0/3, LN#12: 1/2)
<|im_end|>
<|im_start|> pathologist@article{cho2024ie,
title={Extracting lung cancer staging descriptors from pathology reports: a generative language model approach},
author={Hyeongmin Cho, Sooyoung Yoo, Borham Kim, Sowon Jang, Leonard Sunwoo, Sanghwan Kim, Donghyoung Lee, Seok Kim, Sejin Nam, Jin-Haeng Chung},
journal={Journal of Biomedical Informatics},
volume={157},
year={2024},
publisher={Elsevier},
issn={1532-0464},
doi={10.1016/j.jbi.2024.104720},
url={https://doi.org/10.1016/j.jbi.2024.104720}
}