Dataloader/bert_helper.py

import os
import numpy as np
import math
from torch.nn import Tanh, BatchNorm1d
from typing import Optional
import torch.nn as nn
import torch
from transformers import BertModel, BertForSequenceClassification
from transformers import BertTokenizer
from transformers import AutoTokenizer, AutoModel

from torch.utils.data import Dataset as Dataset_n
from torch.utils.data import DataLoader as DataLoader_n
from torch.utils.data import WeightedRandomSampler

def _freeze_bert(
        bert_model: BertModel, freeze_bert=True, freeze_layer_count=-1
):
    """Freeze parameters in BertModel (in place)
    Args:
        bert_model: HuggingFace bert model
        freeze_bert: Bool whether to freeze the bert model
        freeze_layer_count: If freeze_bert, up to what layer to freeze.
    Returns:
        bert_model
    """
    if freeze_bert:
        # freeze the entire bert model
        for param in bert_model.parameters():
            param.requires_grad = False
    else:
        # freeze the embeddings
        for param in bert_model.embeddings.parameters():
            param.requires_grad = False
        if freeze_layer_count != -1:
            if freeze_layer_count > 0 :
                # freeze layers in bert_model.encoder
                for layer in bert_model.encoder.layer[:freeze_layer_count]:
                    for param in layer.parameters():
                        param.requires_grad = False

            if freeze_layer_count < 0 :
                # freeze layers in bert_model.encoder
                for layer in bert_model.encoder.layer[freeze_layer_count:]:
                    for param in layer.parameters():
                        param.requires_grad = False
    return None

def get_frozen_embeder(key_word="bert-large-uncased"):
    tokenizer = AutoTokenizer.from_pretrained(key_word, do_lower_case=False)
    model = AutoModel.from_pretrained(key_word)

    _freeze_bert(model, freeze_bert=True, freeze_layer_count=None)
    return model, tokenizer


def str2emb(string, max_words_num=100, embeder=None, tokenizer=None, reduce_method='mean'):
    string = string.lower()
    str_token = tokenizer(string, return_tensors='pt', max_length=max_words_num,
                          padding='max_length', truncation=True)    
    embeder_output = embeder(**str_token)
    if reduce_method == 'mean':
        embeder_output = torch.mean(embeder_output.last_hidden_state, dim=1)
    elif reduce_method == 'max':
        embeder_output = torch.max(embeder_output.last_hidden_state, dim=1)[0]
    else:
        embeder_output = embeder_output.last_hidden_state
    return embeder_output

def get_synonyms_dict(dict_type=None):
    '''
    Get the dictionary of synonyms for the specified dictionary type
    '''
    if dict_type == 'ROI':
        dict_synonyms = {
            'whole-body': ['whole-body', 'whole body', 'wholebody', 'whole body', 'whole-body', 'whole body', 'wholebody','polytrauma','head-neck-thorax-abdomen-pelvis-leg','head-neck-thorax-abdomen-pelvis'],
            'neck-thorax-abdomen-pelvis-leg': ['neck-thorax-abdomen-pelvis-leg','neck-thx-abd-pelvis-leg', 'angiography neck-thx-abd-pelvis-leg', 'neck thorax abdomen pelvis leg', 'neck and thorax and abdomen and pelvis and leg', 'neck, thorax, abdomen, pelvis & leg', 'neck/thorax/abdomen/pelvis/leg', 'neck, thorax, abdomen, pelvis and leg', 'neck thorax abdomen pelvis leg'],
            'neck-thorax-abdomen-pelvis': ['neck-thorax-abdomen-pelvis', 'neck-thx-abd-pelvis', 'neck thorax abdomen pelvis', 'neck and thorax and abdomen and pelvis', 'neck, thorax, abdomen & pelvis', 'neck/thorax/abdomen/pelvis', 'neck, thorax, abdomen and pelvis', 'neck thorax abdomen & pelvis'],
            'thorax-abdomen-pelvis-leg': ['thorax-abdomen-pelvis-leg','thx-abd-pelvis-leg', 'angiography thx-abd-pelvis-leg', 'thorax abdomen pelvis leg', 'thorax and abdomen and pelvis and leg', 'thorax, abdomen, pelvis & leg', 'thorax/abdomen/pelvis/leg', 'thorax, abdomen, pelvis and leg', 'thorax abdomen pelvis leg'],
            'neck-thorax-abdomen': ['neck-thorax-abdomen', 'neck-thorax-abdomen', 'neck thorax abdomen', 'neck and thorax and abdomen', 'neck, thorax, abdomen', 'neck/thorax/abdomen', 'neck, thorax, abdomen', 'neck thorax abdomen'],
            'head-neck-thorax-abdomen': ['head-neck-thorax-abdomen', 'head-neck-thorax-abdomen', 'head neck thorax abdomen', 'head and neck and thorax and abdomen', 'head, neck, thorax, abdomen', 'head/thorax/abdomen', 'head, thorax, abdomen', 'head thorax abdomen'],
            'head-neck-thorax': ['head-neck-thorax', 'head neck thorax', 'head and neck and thorax', 'head, neck, thorax', 'head/thorax', 'head, thorax', 'head thorax'],
            'thorax-abdomen-pelvis': ['thorax-abdomen-pelvis', 'thx-abd-pelvis', 'polytrauma', 'thorax abdomen pelvis', 'thorax and abdomen and pelvis', 'thorax, abdomen & pelvis', 'thorax/abdomen/pelvis', 'thorax, abdomen and pelvis', 'thorax abdomen & pelvis'],
            'abdomen-pelvis-leg': ['abdomen-pelvis-leg', 'angiography abdomen-pelvis-leg', 'abd-pelvis-leg', 'abdomen pelvis leg', 'abdomen and pelvis and leg', 'abdomen, pelvis & leg', 'abdomen/pelvis/leg', 'abdomen, pelvis, leg', 'abdomen pelvis leg'],
            'neck-thorax': ['neck-thorax', 'neck thorax', 'neck and thorax', 'neck, thorax', 'thorax-neck', 'thorax neck', 'thorax and neck', 'thorax, neck','thorax/neck'],
            'thorax-abdomen': ['thorax-abdomen', 'thorax abdomen', 'thorax and abdomen', 'thorax, abdomen', 'aortic valve'],
            'abdomen-pelvis': ['abdomen-pelvis', 'abdomen pelvis', 'abdomen and pelvis', 'abdomen & pelvis', 'abdomen/pelvis', 'abdomen-pelvis', 'abdomen pelvis', 'abdomen and pelvis', 'abdomen & pelvis', 'abdomen/pelvis'],
            'pelvis-leg': ['pelvis-leg', 'pelvis leg', 'pelvis and leg', 'pelvis, leg', 'pelvis/leg', 'pelvis-leg', 'pelvis leg', 'pelvis and leg', 'pelvis, leg', 'pelvis/leg'],
            'head-neck': ['head-neck', 'head neck', 'head and neck', 'head, neck', 'head/neck', 'head-neck', 'head neck', 'head and neck', 'head, neck', 'head/neck'],
            'abdomen': ['abdomen', 'abdominal', 'belly', 'stomach', 'tummy', 'gut', 'guts', 'viscera', 'bowels', 'intestines', 'gastrointestinal', 'digestive', 'peritoneum','gastric', 'liver', 'spleen', 'pancreas','kidney','lumbar','renal','hepatic','splenic','pancreatic','intervention'],
            'thorax': ['chest', 'thorax', 'breast', 'lung', 'heart','heart-thorakale aorta', 'heart-thorakale', 'mediastinum', 'pleura', 'bronchus', 'bronchi', 'trachea', 'esophagus', 'diaphragm', 'rib', 'sternum', 'clavicle', 'scapula', 'axilla', 'armpit','breast biopsy','thoracic','mammary','caeiothoracic','mediastinal','pleural','bronchial','bronchial tree','tracheal','esophageal','diaphragmatic','costal','sternal','clavicular','scapular','axillary','axillar','cardiac','pericardial','pericardiac','pericardium'],
            'head': ['head', 'headbasis', 'brain', 'skull', 'face','nose','ear','eye','mouth','jaw','cheek','chin','forehead','temporal','parietal','occipital','frontal','mandible','maxilla','mandibular','maxillary','nasal','orbital','orbita','ocular','auricular','otic','oral','buccal','labial','lingual','palatal'],
            'neck': ['neck', 'throat', 'cervical', 'thyroid', 'trachea', 'larynx', 'pharynx', 'esophagus','pharyngeal','laryngeal','cervical','thyroid','trachea','esophagus','carotid','jugular'],
            'hand': ['hand', 'finger', 'thumb', 'palm', 'wrist', 'knuckle', 'fingernail', 'phalanx', 'metacarpal', 'carpal', 'radius'],
            'arm': ['arm', 'forearm', 'upper arm', 'bicep', 'tricep', 'brachium', 'brachial', 'humerus', 'radius', 'ulna', 'elbow', 'shoulder', 'armpit''clavicle', 'scapula', 'acromion', 'acromioclavicular'],
            'leg': ['leg', 'felsenleg','thigh', 'calf', 'shin', 'knee', 'foot', 'ankle', 'toe', 'heel', 'sole', 'arch', 'instep', 'metatarsal', 'phalanx', 'tibia', 'fibula', 'femur', 'patella', 'kneecap','achilles tendon','achilles'],
            'pelvis': ['pelvis', 'hip', 'groin', 'buttock', 'gluteus', 'gluteal', 'ischium', 'pubis', 'sacrum', 'coccyx', 'acetabulum', 'iliac', 'iliac crest', 'iliac spine', 'iliac wing', 'sacroiliac', 'sacroiliac joint', 'sacroiliac ligament', 'sacroiliac spine', 'ureter', 'bladder', 'urethra', 'prostate', 'testicle', 'ovary', 'uterus',],
            'skeleton': ['skeleton','bone','spine', 'back', 'vertebra', 'sacrum', 'coccyx'],
        }
    elif dict_type == 'Label_tissue':
        dict_synonyms = {
            'liver': ['liver','hepatic'],
            'spleen': ['spleen','splenic'],
            'kidney': ['kidney','renal'],
            'pancreas': ['pancreas','pancreatic'],
            'stomach': ['stomach','gastric'],
            'intestine': ['large intestine', 'small intestine','large bowel','small bowel'],
            'gallbladder': ['gallbladder'],
            'adrenal_gland': ['adrenal_gland','adrenal gland'],
            'bladder': ['bladder'],
            'prostate': ['prostate'],
            'uterus': ['uterus'],
            'ovary': ['ovary'],
            'testicle': ['testicle'],
            'lymph_node': ['lymph_node','lymph node'],
            'bone': ['bone'],
            'lung': ['lung'],
            'heart': ['heart'],
            'esophagus': ['esophagus'],
            'muscle': ['muscle'],
            'fat': ['fat'],
            'skin': ['skin'],
            'vessel': ['vessel'],
            'tumor': ['tumor'],
            'other': ['other']
        }
    elif dict_type == 'Task':
        dict_synonyms = {
            'segmentation': ['segmentation', 'seg', 'mask'],
            'classification': ['classification', 'class', 'diagnosis','identify','identification'],
            'localization': ['localization', 'locate', 'location', 'position'],
            'registration': ['registration', 'register', 'align', 'alignment'],
            'detection': ['detection', 'detect', 'find', 'locate'],
            'quantification': ['quantification', 'quantify', 'measure', 'measurement'],
        }
    elif dict_type == 'Modality':
        dict_synonyms = {
            'CT': ['CT', 'computed tomography'],
            'MRI': ['MRI', 'MR', 'magnetic resonance imaging'],
            'PET': ['PET', 'positron emission tomography'],
            'US': ['US', 'ultrasound'],
            'X-ray': ['X-ray', 'radiography'],
            'SPECT': ['SPECT', 'single-photon emission computed tomlogy'],
        }
    else:
        dict_synonyms = {
            '\'gender\'': ['\'gender\'', '\'sex\'', '\'M/F\'', '\'m/f\''],
            '\'modality\'': ['\'modality\'', '\'modal\''],
            '\'male\'': ['\'male\'', '\'m\''],
            '\'female\'': ['\'female\'', '\'f\'','\'woman\''],
            '\'high-grade glioma\'': ['\'high-grade glioma\'', '\'high grade glioma\'', '\'HGG\''],
            '\'low-grade glioma\'': ['\'low-grade glioma\'', '\'low grade glioma\'', '\'LGG\''],
            '\'atlas scaling factor\'': ['\'atlas scaling factor\'', '\'asf\''],
            '\'age\'': ['\'age\'', '\'years\'', '\'year\'', '\'y/o\'', '\'y.o.\''],
            '\'education\'': ['\'educ\'', '\'educat\'', '\'education\''],
            '\'roi\'': ['\'roi\'', '\'region of interest\'', '\'region\''],
            '\'mini-mental state examination\'': ['\'mini-mental state examination\'', '\'mmse\''],
            '\'clinical dementia rating\'': ['\'clinical dementia rating\'', '\'cdr\''],
            '\'socio-economic status\'': ['\'socio-economic status\'', '\'ses\''],
            '\'unknown\'': ['\'unknown\'', '\'unkn\'', '\'not available\'', '\'nan\'', '\'n/a\'', '\'none\'', '\'n.a.\'', '\'not applicable\'','\'not specified\'', '\'unspecified\'', '\'not given\'', '\'null\''],
            '': [' segmentation', '\'seg\'', '\'registration\''],
        }
    return dict_synonyms

def replace_text(text, dict_synonyms):
    '''
    Replace the text in the text with the standard term
    '''
    if isinstance(text, str):
        for key, value in dict_synonyms.items():
            for v in value:
                if v.lower() in text.lower():
                    text = text.replace(v, key)
        return text
    elif isinstance(text, list):
        text = [replace_text(t, dict_synonyms) for t in text]
    elif isinstance(text, dict):
        for key in text.keys():
            # replace values in dict
            text[key] = replace_text(text[key], dict_synonyms)
            # replace keys in dict
            for k in dict_synonyms.keys():
                if k.lower() in key.lower():
                    text[dict_synonyms[k]] = text.pop(key)
    return text


def replace_synonyms(text, dict_synonyms):
    '''
    Replace the synonyms in the text with the standard term
    '''
    if isinstance(text,str):
        for key, value in dict_synonyms.items():
            for v in value:
                if v.lower() in text.lower():
                    return key
        Warning(f"Value {text} is not in the correct format")
    elif isinstance(text,list):
        text = [replace_synonyms(t, dict_synonyms) for t in text]
    elif isinstance(text,dict):
        for key in text.keys():
            # replace values in dict
            text[key] = replace_synonyms(text[key], dict_synonyms)
            # replace keys in dict
            for k in dict_synonyms.keys():
                text[dict_synonyms[k]] = text.pop(key)
    return text    

if __name__ == "__main__":
    # model_name = "bert-base-uncased"
    # model_name = "bert-large-uncased"
    model_name = "/home/jachin/data/Github/OmniMorph/External/Models/bert_large_uncased"
    # model_name = "Rostlab/prot_bert"
    # model_name = "fspanda/Medical-Bio-BERT2"
    # model_name = "GerMedBERT/medbert-512"

    reduce_method = 'mean'
    max_words_num = 32  # max number of words in the caption > 2
    
    embeder, tokenizer = get_frozen_embeder(model_name)

    # string1 = ["mri", "female"]
    string1 = "modality: ct, gender: female, age: 51, roi: abdomen"
    # string1 = "modality: Magnetic Resonance, gender: female"
    embeder_output1 = str2emb(string1, max_words_num, embeder, tokenizer, reduce_method=reduce_method)

    # string2 = "Hello world!"
    # string2 = ["ct", "male"]
    # string2 = "modality: mri, gender: female, roi: head"
    string2 = "modality: ct, gender: female, age: 50, roi: head"
    # string2 = "modality: ct, gender: male, roi: head"
    embeder_output2 = str2emb(string2, max_words_num, embeder, tokenizer, reduce_method=reduce_method)

    input_size = embeder.config.vocab_size
    in_size = embeder.config.hidden_size
    
    print(embeder, input_size, in_size)
    print(tokenizer)
    
    # embeder_output1 shape: [batch_size, max_words_num, hidden_size]
    print(embeder_output1)
    print(embeder_output1.shape)  # torch.Size([1, 8, 768])
    
    # embeder_output2 shape: [batch_size, max_words_num, hidden_size]
    print(embeder_output2)
    print(embeder_output2.shape)  # torch.Size([1, 8, 768])
    
    # check the difference between the two sentences in embedding space
    # embeder_output1[0, :, :] shape: [max_words_num, hidden_size]
    # embeder_output2[0, :, :] shape: [max_words_num, hidden_size]
    # error = torch.max(torch.abs(embeder_output1[0, :, :] - embeder_output2[0, :, :]), dim=-1)
    error = torch.abs(embeder_output1 - embeder_output2)
    print(error)
    print("Embedding distance between the two sentences: ")
    print(f"String1: {string1}")
    print(f"String2: {string2}")
    print(torch.mean(error))
    exit()