feat(intl/eu): 75x speedup text data formatting

Author: zietzm

onsides_intl/onsides_eu/pyproject.toml

@@ -4,7 +4,9 @@ version = "0.1.0"
 description = "ONSIDES European Union"
 readme = "README.md"
 requires-python = ">=3.11"
-dependencies = []
+dependencies = [
+    "pyahocorasick>=2.1.0",
+]
 
 [build-system]
 requires = ["hatchling"]

onsides_intl/onsides_eu/src/onsides_eu/stringsearch.py

@@ -0,0 +1,66 @@
+import ahocorasick
+from pydantic import BaseModel
+
+
+class MeddraSearchTerm(BaseModel):
+    term: str
+    meddra_pt_code: int
+
+
+class FoundMeddraTerm(BaseModel):
+    term: str
+    meddra_pt_code: int
+    start: int
+    end: int
+
+
+def build_meddra_search_tree(
+    meddra_terms: list[MeddraSearchTerm],
+) -> ahocorasick.Automaton:
+    """
+    Builds an Aho-Corasick tree from a list of MedDRA terms.
+    """
+    tree = ahocorasick.Automaton(str, str)
+    for meddra_obj in meddra_terms:
+        tree.add_word(meddra_obj.term, meddra_obj.model_dump_json())
+    tree.make_automaton()
+    return tree
+
+
+def find_meddra_terms_in_text(
+    text: str,
+    meddra_tree: ahocorasick.Automaton,
+) -> list[FoundMeddraTerm]:
+    """
+    Finds all MedDRA terms in a text using an Aho-Corasick tree.
+    """
+    found_terms = list()
+    for end_index, obj_json in meddra_tree.iter(text):
+        meddra_obj = MeddraSearchTerm.model_validate_json(obj_json)
+        start_index = end_index - len(meddra_obj.term) + 1
+        obj = FoundMeddraTerm(
+            term=meddra_obj.term,
+            meddra_pt_code=meddra_obj.meddra_pt_code,
+            start=start_index,
+            end=end_index,
+        )
+        found_terms.append(obj)
+    return found_terms
+
+
+def build_bert_string(
+    text: str,
+    match: FoundMeddraTerm,
+    nwords: int = 125,
+    prop_before: float = 0.125,
+) -> str:
+    term_nwords = len(match.term.split())
+    n_words_before = prop_before * (nwords - 2 * term_nwords)
+    n_words_after = (1 - prop_before) * (nwords - 2 * term_nwords)
+    n_words_before = max(int(n_words_before), 1)
+    n_words_after = max(int(n_words_after), 1)
+    before_words = text[: match.start].split()[-n_words_before:]
+    after_words = text[match.end :].split()[:n_words_after]
+    words_list = [match.term] + before_words + ["EVENT"] + after_words
+    result = " ".join(words_list)
+    return result

onsides_intl/onsides_eu/src/onsides_eu/text_data_format.py

@@ -1,114 +1,120 @@
-import numpy as np 
-import pandas as pd
-import requests
-from tqdm import tqdm
-from glob import glob
-import ast, re, json, orjson
-from time import sleep
+"""Find MedDRA term exact matches in drug label free text. Setup data for the
+OnSIDES model.
+
+There are a couple things that I did here which should be clarified.
+
+First, I don't do any RxNorm mapping here. I saw that previous code attempted to
+do this, but it was trying to join drug names from EMA to RxNorm SET IDs, which
+look like UUIDs. Text joins didn't work, so I just explicitly set those columns
+to None below, just to ensure we have all the same columns as the original code.
+
+Second, I'm not sure about the "AR" section. That's what the previous code did,
+so I just did it here too.
+
+Third, I used the same MedDRA terms as previous code (I think). Not sure why
+we're only using 5 character or longer terms. My code looks for either PT or LLT
+and maps to PT terms.
+"""
+
 import argparse
-import warnings
-warnings.filterwarnings('ignore')
-import os
+import logging
+import pathlib
+
+import polars as pl
+import tqdm.auto as tqdm
+
+from onsides_eu.stringsearch import (
+    MeddraSearchTerm,
+    build_bert_string,
+    build_meddra_search_tree,
+    find_meddra_terms_in_text,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def format_text(
+    data_folder: pathlib.Path,
+    external_data_folder: pathlib.Path,
+) -> None:
+    drug_to_ade_text = (
+        pl.read_csv(data_folder / "ade_text_table.csv")
+        .with_columns(pl.col("ade_text").str.to_lowercase())
+        .select("drug", "ade_text")
+        .to_dicts()
+    )
+    meddra_df = (
+        pl.read_csv(external_data_folder / "umls_meddra_en.csv")
+        .filter(
+            pl.col("TTY").is_in({"PT", "LLT"}),
+        )
+        .with_columns(
+            pl.col("STR").str.to_lowercase().alias("term"),
+        )
+        .rename({"SDUI": "meddra_pt_code"})
+    )
+    meddra_pt_code_to_term = (
+        meddra_df.filter(pl.col("TTY").eq("PT"))
+        .select("STR", "meddra_pt_code")
+        .to_pandas()
+        .set_index("meddra_pt_code")["STR"]
+        .to_dict()
+    )
+    meddra_terms = (
+        meddra_df.filter(pl.col("term").str.len_chars().ge(5))
+        .select("term", "meddra_pt_code")
+        .unique()
+        .to_dicts()
+    )
+    meddra_terms = [MeddraSearchTerm.model_validate(t) for t in meddra_terms]
+    logger.info(
+        f"Found {len(drug_to_ade_text)} drugs. "
+        f"Searching for exact matches of {len(meddra_terms)} MedDRA terms."
+    )
+    meddra_tree = build_meddra_search_tree(meddra_terms)
+
+    exact_terms = list()
+    for drug_term in tqdm.tqdm(drug_to_ade_text):
+        ade_text = drug_term["ade_text"]
+        matches = find_meddra_terms_in_text(ade_text, meddra_tree)
+        for match in matches:
+            bert_string = build_bert_string(ade_text, match)
+            row = {
+                "label_id": drug_term["drug"],
+                "found_term": match.term,
+                "location": match.start,
+                "string": bert_string,
+                "section": "AR",
+                "set_id": drug_term["drug"],
+                "drug": None,
+                "spl_version": None,
+                "pt_meddra_id": match.meddra_pt_code,
+                "pt_meddra_term": meddra_pt_code_to_term.get(match.meddra_pt_code),
+            }
+            exact_terms.append(row)
+
+    logger.info(f"Found {len(exact_terms)} exact matches.")
+    pl.DataFrame(exact_terms).write_csv(data_folder / "bert_input_v2.csv")
+
 
 def main():
-    parser = argparse.ArgumentParser(description='let the code know where the data is held')
-    parser.add_argument('--data_folder', required=True, help='Path to the data folder.')
-    parser.add_argument('--external_data', required=True, help='Path to the where the external data is housed.')
-    parser.add_argument('--map_folder', required=True, help='Path to the where the external data used for OnSIDES model is housed.')
+    logging.basicConfig(level=logging.INFO)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data_folder",
+        type=pathlib.Path,
+        required=True,
+        help="Path to the data folder.",
+    )
+    parser.add_argument(
+        "--external_data",
+        type=pathlib.Path,
+        required=True,
+        help="Path to the external data folder.",
+    )
     args = parser.parse_args()
-    data_folder = args.data_folder
-    external_data_folder = args.external_data
-    map_folder = args.map_folder
-
-    #read in table for drug-ade free-text data
-    ade_text_table_df = pd.read_csv(data_folder+'ade_text_table.csv')
-
-    ##Standard Vocabulary Mapping - here, we will use the UMLS MedDRA tables.
-    meddra_df = pd.read_csv(external_data_folder+'umls_meddra_en.csv')
-    meddra_df['STR'] = meddra_df.STR.apply(lambda x: x.lower())
-    meddra_df['len'] = meddra_df.STR.apply(lambda x: len(x))
-    meddra_dict = dict(zip(meddra_df.STR, meddra_df.SDUI))
-    meddra_df = meddra_df[(meddra_df.TTY == 'PT')|(meddra_df['len'] > 5)]
-
-    exact_terms = []
-    for i, row in tqdm(ade_text_table_df.iterrows()):
-        label_id = row['drug']
-        text = row['ade_txt'].lower()
-        found_terms = list()
-        for mdr_term in meddra_dict.keys():
-            if text.find(mdr_term) == -1:
-                continue
-            else:
-                li = text.split(mdr_term)
-                start_pos = 0
-                for i in range(len(li)-1):
-                    # the occurrence of the word is at the end of the previous string
-                    start_pos = sum([len(li[j]) for j in range(i+1)]) + i*len(mdr_term)
-                    if not mdr_term == text[start_pos:(start_pos+len(mdr_term))]:
-                        raise Exception(f" mdr_term: '{mdr_term}', term_in_text: '{text[start_pos:(start_pos+len(mdr_term))]}'")
-                    found_terms.append((mdr_term, meddra_dict[mdr_term], start_pos, len(mdr_term)))
-        exact_terms.append([label_id, found_terms])
-
-    exact_terms_df = pd.DataFrame(exact_terms, columns=['label_id', 'found_terms'])
-    exact_terms_df = exact_terms_df.explode('found_terms')
-    exact_terms_df['len'] = exact_terms_df['found_terms'].apply(lambda x: x[3] if str(x) != 'nan' else None)
-    exact_terms_df = exact_terms_df[exact_terms_df['len'] >= 5]
-    exact_terms_df['found_term'] = exact_terms_df['found_terms'].apply(lambda x: x[0] if str(x) != 'nan' else None)
-    exact_terms_df['meddra_id'] = exact_terms_df['found_terms'].apply(lambda x: x[1] if str(x) != 'nan' else None)
-    exact_terms_df['location'] = exact_terms_df['found_terms'].apply(lambda x: x[2] if str(x) != 'nan' else None)
-    exact_terms_df = exact_terms_df.drop(['found_terms', 'len'], axis = 1)
-
-    building_strings = []
-    ade_text_table_dict = dict(zip(ade_text_table_df.drug, ade_text_table_df.ade_txt))
-    for i, row in tqdm(exact_terms_df.iterrows()):
-        term, label_id, start_pos = row['found_term'], row['label_id'], row['location']
-        #default settings
-        nwords, prop_before = 125, 0.125
-        #pull the full text
-        ar_text = ade_text_table_dict[label_id]
-
-        term_nwords = len(term.split())
-        size_before = max(int((nwords-2*term_nwords)*prop_before), 1)
-        size_after = max(int((nwords-2*term_nwords)*(1-prop_before)), 1)
-
-        before_text = ar_text[:start_pos]
-        after_text = ar_text[(start_pos+term_nwords):]
-
-        before_parts = before_text.split()[-1*size_before:]
-        after_parts = after_text.split()[:size_after]
-
-        li = [term]
-        li.extend(before_parts)
-        li.append('EVENT')
-        li.extend(after_parts)
-        example_string = ' '.join(li)
-        building_strings.append(example_string)
-        exact_terms_df['string'] = building_strings
-    
-    #save dataframe
-    exact_terms_df.to_csv(data_folder+'sentences-rx_method14_nwords125_clinical_bert_application_set_AR.csv', index=False)
-
-    #further prep the data for the model
-    #required columns : section, drug, label_id, set_id, spl_version, pt_meddra_id, pt_meddra_term
-    exact_terms_df = pd.read_csv(data_folder+'sentences-rx_method14_nwords125_clinical_bert_application_set_AR_v0924.csv')
-    exact_terms_df['section'] = 'AR'
-    exact_terms_df['set_id'] = exact_terms_df['label_id']
-    
-    drug_map = pd.read_csv(map_folder+'spl/maps/20230512/rxnorm_mappings.txt', delimiter = '|')
-    drug_id_dict = dict(zip(drug_map.SETID, drug_map.RXCUI))
-    drug_ver_dict = dict(zip(drug_map.SETID, drug_map.SPL_VERSION))
-    exact_terms_df['drug'] = exact_terms_df.set_id.apply(lambda x: drug_id_dict[x] if x in drug_id_dict.keys() else None)
-    exact_terms_df['spl_version'] = exact_terms_df.set_id.apply(lambda x: drug_ver_dict[x] if x in drug_ver_dict.keys() else None)
-
-    llt_pt = pd.read_csv(map_folder+'meddra_llt_pt_map.txt', delimiter = '|')
-    llt_pt_id_dict = dict(zip(llt_pt.llt_concept_code, llt_pt.pt_concept_code))
-    llt_pt_term_dict = dict(zip(llt_pt.llt_concept_code, llt_pt.pt_concept_name))
-    exact_terms_df['pt_meddra_id'] = exact_terms_df.meddra_id.apply(lambda x: llt_pt_id_dict[x] if x in llt_pt_id_dict.keys() else None)
-    exact_terms_df['pt_meddra_term'] =  exact_terms_df.meddra_id.apply(lambda x: llt_pt_term_dict[x] if x in llt_pt_term_dict.keys() else None)
-
-    #save dataframe
-    exact_terms_df.to_csv(data_folder+'sentences-rx_method14_nwords125_clinical_bert_application_set_AR.csv', index=False)
-
-
-if __name__ == '__main__':
-    main()
+    format_text(args.data_folder, args.external_data)
+
+
+if __name__ == "__main__":
+    main()