Downloads: 1
Switch to side-by-side view
--- a +++ b/utils.py @@ -0,0 +1,677 @@ +from typing import List, Tuple, Callable, Dict, Union, Iterable +from annotations import Entity, Relation +from ehr import HealthRecord + +import os +import sys +from pickle import dump, load +from IPython.core.display import display, HTML +import json +from collections import defaultdict +import pandas as pd +import networkx as nx +import math +import matplotlib.pyplot as plt +from io import BytesIO +import base64 +import matplotlib + + +TPL_HTML = """<span style = "background-color: {color}; border-radius: 5px;"> {content} </span>""" + +TPL_HTML_HOVER = """<span style = "background-color: {color}; border-radius: 5px;" class="{grp}"> {content} <span style = "background: {color}">{ent_type}</span></span>""" + +COLORS = {"Drug": "#aa9cfc", "Strength": "#ff9561", + "Form": "#7aecec", "Frequency": "#9cc9cc", + "Route": "#ffeb80", "Dosage": "#bfe1d9", + "Reason": "#e4e7d2", "ADE": "#ff8197", + "Duration": "#97c4f5"} + + +def add_ent_group(entities: Union[Dict[str, Entity], List[Entity]], + relations: Union[Dict[str, Relation], List[Relation]]) -> List[Entity]: + """ + Adds relation group to Entity objects. + + Parameters + ---------- + entities : Union[Dict[str, Entity], List[Entity]] + Entities + + relations : Union[Dict[str, Relation], List[Relation]]) + Relations + + Returns + ------- + List[Entity] + List of Entities with group information added. + """ + + # Convert entities to a dictionary if not + if not isinstance(entities, dict): + ent_dict = {} + for ent in entities: + ent_dict[ent.ann_id] = ent + entities = ent_dict + + # Add group information + for rel in relations: + entities[rel.arg1.ann_id].relation_group += "group-" + rel.ann_id + " " + entities[rel.arg2.ann_id].relation_group += "group-" + rel.ann_id + " " + + return list(entities.values()) + + +# noinspection PyTypeChecker +def display_ehr(text: str, + entities: Union[Dict[str, Entity], List[Entity]], + relations: Union[Dict[str, Relation], List[Relation]] = None, + return_html: bool = False) -> Union[None, str]: + """ + Highlights EHR records with colors and displays + them as HTML. Ideal for working with Jupyter Notebooks + + Parameters + ---------- + text : str + EHR record to render + + entities : Union[Dict[str, Entity], List[Entity]] + A list of Entity objects + + relations : Union[Dict[str, Relation], List[Relation]] + A list of relations. If provided, entities should be a dictionary. + + return_html : bool + Indicator for returning HTML or printing the tagged EHR. + The default is False. + + Returns + ------- + Union[None, str] + If return_html is true, returns html strings + otherwise displays HTML. + + """ + if relations is not None: + entities = add_ent_group(entities, relations) + + if isinstance(entities, dict): + entities = list(entities.values()) + + # Sort entity by starting range + entities.sort(key=lambda x: x.range[0]) + + # Final text to render + render_text = "" + start_idx = 0 + + # Display legend + if not return_html: + for ent, col in COLORS.items(): + render_text += TPL_HTML.format(content=ent, color=col) + render_text += " " * 5 + + render_text += '\n' + render_text += '--' * 50 + render_text += "\n\n" + + # Replace each character range with HTML span template + for ent in entities: + if start_idx > ent.range[0]: + continue + + render_text += text[start_idx:ent.range[0]] + + if return_html: + render_text += TPL_HTML_HOVER.format( + content=text[ent.range[0]:ent.range[1]], + color=COLORS[ent.name], + grp=ent.relation_group, + ent_type=ent.name) + else: + render_text += TPL_HTML.format( + content=text[ent.range[0]:ent.range[1]], + color=COLORS[ent.name]) + + start_idx = ent.range[1] + + render_text += text[start_idx:] + render_text = render_text.replace("\n", "<br>") + + if return_html: + return render_text + else: + display(HTML(render_text)) + + +def display_knowledge_graph(long_relation_df: pd.DataFrame, num_col: int = 2, + height: int = 8, width: int = 8, + return_html: bool = False) -> Union[None, str]: + """ + Highlights EHR records with colors and displays + them as HTML. Ideal for working with Jupyter Notebooks + + Parameters + ---------- + long_relation_df: pd.DataFrame + Relation dataframe in long format. Should have columns named: + ['drug_id', 'drug', 'arg', 'edge'] + + num_col: int + Number of columns in the grid. Number of rows are automatically + calculated based on this. The default is 2. + + height: int + The height of a single graph in inches. The default is 6. + + width: int + The width of a single graph in inches. The default is 6. + + return_html: bool + Indicator for returning the HTML img tag or displaying the plot. + The default is False. + + Returns + ------- + Union[None, str] + If return_html is true, returns html string + otherwise displays the plot. + + """ + if return_html: + matplotlib.use('Agg') + + drug_ids = sorted(list(pd.unique(long_relation_df['drug_id']))) + num_row = math.ceil(len(drug_ids) / num_col) + + if num_row == 0: + return None + + _ = plt.subplots(num_row, num_col, figsize=(num_col * width, height * num_row)) + + i = 0 + for i, d in enumerate(drug_ids): + sub_rel = long_relation_df[long_relation_df["drug_id"] == d] + labels = sub_rel.set_index(['drug', 'arg'])['edge'].to_dict() + + plt.subplot(num_row, num_col, i + 1) + + # Knowledge graph for a single drug + graph = nx.from_pandas_edgelist(sub_rel, "drug", "arg", edge_attr=True, create_using=nx.MultiDiGraph()) + + # Drug will always be the first in the graph + color_map = ['#aa9cfc'] + ['skyblue'] * (len(graph.nodes) - 1) + + pos = nx.spring_layout(graph) + + # Draw the graph + nx.draw(graph, with_labels=True, font_size=12, pos=pos, + node_color=color_map, node_size=2000) + + # Draw edge labels + nx.draw_networkx_edge_labels(graph, edge_labels=labels, + pos=pos, font_color='red') + + # Remove axis for empty plots, if any + i += 1 + while i < num_row * num_col: + plt.subplot(num_row, num_col, i + 1) + plt.axis('off') + i += 1 + + if not return_html: + plt.show() + return + + # Create an encoding for the image + tmp_file = BytesIO() + + plt.tight_layout() + plt.savefig(tmp_file, format="png") + + encoded = base64.b64encode(tmp_file.getvalue()).decode('utf-8') + img_tag = '<img id="knowledge-graph" src=\'data:image/png;base64,{}\'>'.format(encoded) + + return img_tag + + +def read_data(data_dir: str = 'data/', + tokenizer: Callable[[str], List[str]] = None, + is_bert_tokenizer: bool = True, + verbose: int = 0) -> Tuple[List[HealthRecord], List[HealthRecord]]: + """ + Reads train and test data + + Parameters + ---------- + data_dir : str, optional + Directory where the data is located. + It should have directories named 'train' and 'test' + The default is 'data/'. + + tokenizer : Callable[[str], List[str]], optional + The tokenizer function to use.. The default is None. + + is_bert_tokenizer : bool + If the tokenizer is a BERT-based WordPiece tokenizer + + verbose : int, optional + 1 to print reading progress, 0 otherwise. The default is 0. + + Returns + ------- + Tuple[List[HealthRecord], List[HealthRecord]] + Train data, Test data. + + """ + train_path = os.path.join(data_dir, "train") + test_path = os.path.join(data_dir, "test") + + # Get all IDs for train and test data + train_ids = list(set(['.'.join(fname.split('.')[:-1]) \ + for fname in os.listdir(train_path) \ + if not fname.startswith('.')])) + + test_ids = list(set(['.'.join(fname.split('.')[:-1]) \ + for fname in os.listdir(test_path) \ + if not fname.startswith('.')])) + + if verbose == 1: + print("Train data:") + + train_data = [] + for idx, fid in enumerate(train_ids): + record = HealthRecord(fid, text_path=os.path.join(train_path, fid + '.txt'), + ann_path=os.path.join(train_path, fid + '.ann'), + tokenizer=tokenizer, + is_bert_tokenizer=is_bert_tokenizer) + train_data.append(record) + if verbose == 1: + draw_progress_bar(idx + 1, len(train_ids)) + + if verbose == 1: + print('\n\nTest Data:') + + test_data = [] + for idx, fid in enumerate(test_ids): + record = HealthRecord(fid, text_path=os.path.join(test_path, fid + '.txt'), + ann_path=os.path.join(test_path, fid + '.ann'), + tokenizer=tokenizer, + is_bert_tokenizer=is_bert_tokenizer) + test_data.append(record) + if verbose == 1: + draw_progress_bar(idx + 1, len(test_ids)) + + return train_data, test_data + + +def read_ade_data(ade_data_dir: str = 'ade_data/', + verbose: int = 0) -> List[Dict]: + """ + Reads train and test ADE data + + Parameters + ---------- + + ade_data_dir : str, optional + Directory where the ADE data is located. The default is 'ade_data/'. + + verbose : int, optional + 1 to print reading progress, 0 otherwise. The default is 0. + + Returns + ------- + List[Dict] + ADE data + + """ + + # Get all the IDs of ADE data + ade_file_ids = sorted(list(set(['.'.join(fname.split('.')[:-1]) \ + for fname in os.listdir(ade_data_dir) \ + if not fname.startswith('.')]))) + + # Load ADE data + ade_data = [] + for idx, fid in enumerate(ade_file_ids): + with open(ade_data_dir + fid + '.json') as f: + data = json.load(f) + ade_data.extend(data) + + ade_data = process_ade_files(ade_data) + if verbose == 1: + print("\n\nADE data: Done") + + return ade_data + + +def process_ade_files(ade_data: List[dict]) -> List[dict]: + """ + Extracts tokens and creates Entity and Relation objects + from raw json data. + + Parameters + ---------- + ade_data : List[dict] + Raw json data. + + Returns + ------- + List[dict] + Tokens, entities and relations. + + """ + ade_records = [] + + for ade in ade_data: + entities = {} + relations = {} + relation_backlog = [] + + # Tokens + tokens = ade['tokens'] + + # Entities + e_num = 1 + for ent in ade['entities']: + ent_id = 'T' + "%s" % e_num + if ent['type'] == 'Adverse-Effect': + ent['type'] = 'ADE' + + ent_obj = Entity(entity_id=ent_id, + entity_type=ent['type']) + + r = [ent['start'], ent['end'] - 1] + r = list(map(int, r)) + ent_obj.set_range(r) + + text = '' + for token_ent in ade['tokens'][ent['start']:ent['end']]: + text += token_ent + ' ' + ent_obj.set_text(text) + + entities[ent_id] = ent_obj + e_num += 1 + + # Relations + r_num = 1 + for relation in ade['relations']: + rel_id = 'R' + "%s" % r_num + rel_details = 'ADE-Drug' + entity1 = "T" + str(relation['head'] + 1) + entity2 = "T" + str(relation['tail'] + 1) + + if entity1 in entities and entity2 in entities: + rel = Relation(relation_id=rel_id, + relation_type=rel_details, + arg1=entities[entity1], + arg2=entities[entity2]) + + relations[rel_id] = rel + + else: + relation_backlog.append([rel_id, rel_details, + entity1, entity2]) + r_num += 1 + + ade_records.append({"tokens": tokens, "entities": entities, "relations": relations}) + return ade_records + + +def map_entities(entities: Union[Dict[str, Entity], List[Entity]], + actual_relations: Union[Dict[str, Relation], List[Relation]] = None) \ + -> Union[List[Tuple[Relation, None]], List[Tuple[Relation, int]]]: + """ + Maps each drug entity to all other non-drug entities in the list. + + Parameters + ---------- + entities : List[Entity] + List of entities. + + actual_relations : List[Relation], optional + List of actual relations (for training data). + The default is None. + + Returns + ------- + Union[List[Relations], List[Tuple[Relation, int]]] + List of mapped relations. If actual relations are specified, + also returns a flag to indicate if it is an actual relation. + + """ + + drug_entities = [] + non_drug_entities = [] + + if isinstance(entities, dict): + entities = list(entities.values()) + + if actual_relations and isinstance(actual_relations, dict): + actual_relations = list(actual_relations.values()) + + # Splitting each entity to drug and non-drug entities + for ent in entities: + if ent.name.lower() == "drug": + drug_entities.append(ent) + else: + non_drug_entities.append(ent) + + relations = [] + i = 1 + + # Mapping each drug entity to each non-drug entity + for ent1 in drug_entities: + for ent2 in non_drug_entities: + rel = Relation(relation_id="R%d" % i, + relation_type=ent2.name + "-Drug", + arg1=ent1, arg2=ent2) + relations.append(rel) + i += 1 + + if actual_relations is None: + return list(zip(relations, [None] * len(relations))) + + # Maps each relation type to list of actual relations + actual_rel_dict = defaultdict(list) + for rel in actual_relations: + actual_rel_dict[rel.name].append(rel) + + relation_flags = [] + flag = 0 + + # Computes actual relation flags + for rel in relations: + for act_rel in actual_rel_dict[rel.name]: + if rel == act_rel: + flag = 1 + break + + relation_flags.append(flag) + flag = 0 + + return list(zip(relations, relation_flags)) + + +def get_long_relation_table(relations: Iterable[Relation]) -> pd.DataFrame: + """ + Returns the relations in a long table format with the columns + ['drug_id', 'drug', 'arg', 'edge'] where arg is entity related + to drug and edge is the entity type. + + Parameters + ---------- + relations : Iterable[Relation] + A list of relations. + + Returns + ------- + pd.DataFrame + All the relations in a long tabular format. + + """ + rel_dict = {'drug_id': [], 'drug': [], 'arg': [], 'edge': []} + + for rel in relations: + if rel.arg1.name == "Drug": + rel_dict['drug_id'].append(rel.arg1.ann_id) + rel_dict['drug'].append(rel.arg1.ann_text) + rel_dict['arg'].append(rel.arg2.ann_text) + + else: + rel_dict['drug_id'].append(rel.arg2.ann_id) + rel_dict['drug'].append(rel.arg2.ann_text) + rel_dict['arg'].append(rel.arg1.ann_text) + + rel_dict['edge'].append(rel.name.split('-')[0]) + + rel_df = pd.DataFrame(rel_dict) + return rel_df + + +def get_relation_table(relations: Union[pd.DataFrame, Iterable[Relation]], + is_long_df: bool = True) -> pd.DataFrame: + """ + Returns the relations in a wide table format. + + Parameters + ---------- + relations : Union[pd.DataFrame, Iterable[Relation]] + Either a list of relations, or relations table in long format. + + is_long_df : bool + Indicator for relations parameter. True indicates the input is + a long dataframe. False indicates it is a list of relations. + + Returns + ------- + str + HTML blob of all the relations in a tabular format. + + """ + relations = relations.drop_duplicates() + + if not is_long_df: + relations = get_long_relation_table(relations) + + relations = relations.rename(columns={"drug_id": "Drug ID", "drug": "Drug", + "edge": "Entity Type", "arg": "Entity Text"}) + + relation_df = ( + relations + .groupby(["Drug ID", "Drug", "Entity Type"])["Entity Text"] + .apply(lambda x: list(x)) + .reset_index(name="Entity Text") + .set_index(["Drug ID", "Drug", "Entity Type"]) + ) + + relation_df["Entity Text"] = relation_df["Entity Text"].apply(lambda x: "\n".join(x)) + + empty_header = " <tr style=\"text-align: right;\">\n <th></th>\n <th></th>\n <th></th>\n <th>Entity Text</th>\n </tr>\n" + empty_colname = "<th></th>" + + relation_html = ( + relation_df + .to_html(classes=['table'], border=0) + .replace("\\n", "<br>") + .replace(empty_header, "") + .replace(empty_colname, "<th>Entity Text</th>") + ) + return relation_html + + +def draw_progress_bar(current, total, string='', bar_len=20): + """ + Draws a progress bar, like [====> ] 40% + + Parameters + ------------ + current: int/float + Current progress + + total: int/float + The total from which the current progress is made + + string: str + Additional details to write along with progress + + bar_len: int + Length of progress bar + """ + percent = current / total + arrow = ">" + if percent == 1: + arrow = "" + # Carriage return, returns to the beginning of line to overwrite + sys.stdout.write("\r") + sys.stdout.write("Progress: [{:<{}}] {}/{}".format("=" * int(bar_len * percent) + arrow, + bar_len, current, total) + string) + sys.stdout.flush() + + +def is_whitespace(char): + """ + Checks if the character is a whitespace + + Parameters + -------------- + char: str + A single character string to check + """ + # ord() returns unicode and 0x202F is the unicode for whitespace + if char == " " or char == "\t" or char == "\r" or char == "\n" or ord(char) == 0x202F: + return True + else: + return False + + +def is_punct(char): + """ + Checks if the character is a punctuation + + Parameters + -------------- + char: str + A single character string to check + """ + if char == "." or char == "," or char == "!" or char == "?" or char == '\\': + return True + else: + return False + + +def save_pickle(file, variable): + """ + Saves variable as a pickle file + + Parameters + ----------- + file: str + File name/path in which the variable is to be stored + + variable: object + The variable to be stored in a file + """ + if file.split('.')[-1] != "pkl": + file += ".pkl" + + with open(file, 'wb') as f: + dump(variable, f) + print("Variable successfully saved in " + file) + + +def open_pickle(file): + """ + Returns the variable after reading it from a pickle file + + Parameters + ----------- + file: str + File name/path from which variable is to be loaded + """ + if file.split('.')[-1] != "pkl": + file += ".pkl" + + with open(file, 'rb') as f: + return load(f)