ddi-fw 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

ddi_fw/datasets/__init__.py

@@ -0,0 +1,12 @@
+from .core import BaseDataset
+from .ddi_mdl.base import DDIMDLDataset 
+from .mdf_sa_ddi.base import MDFSADDIDataset
+from .custom import CustomDataset
+from .embedding_generator import create_embeddings
+from .embedding_generator_new import EmbeddingGenerator,PretrainedEmbeddingGenerator,SBertEmbeddingGenerator,LLMEmbeddingGenerator,create_embeddings_new
+from .idf_helper import IDF
+from .feature_vector_generation import SimilarityMatrixGenerator, VectorGenerator
+
+__all__  = ['BaseDataset','DDIMDLDataset','MDFSADDIDataset']
+
+

ddi_fw/datasets/core.py

@@ -0,0 +1,416 @@
+import glob
+from sklearn.model_selection import KFold, StratifiedKFold, train_test_split
+from sklearn.preprocessing import LabelBinarizer
+from abc import ABC, abstractmethod
+import numpy as np
+import pandas as pd
+import pathlib
+from datasets.idf_helper import IDF
+
+from utils.zip_helper import ZipHelper
+from .feature_vector_generation import SimilarityMatrixGenerator, VectorGenerator
+from ner.ner import CTakesNER
+from utils import create_folder_if_not_exists
+from stopwatch import Stopwatch, profile
+
+HERE = pathlib.Path(__file__).resolve().parent
+
+
+def stack(df_column):
+    return np.stack(df_column.values)
+
+
+class BaseDataset(ABC):
+    def __init__(self, chemical_property_columns, embedding_columns, ner_columns,
+                 **kwargs):
+        self.__similarity_related_columns__ = []
+        self.__similarity_related_columns__.extend(chemical_property_columns)
+        self.__similarity_related_columns__.extend(ner_columns)
+
+        self.chemical_property_columns = chemical_property_columns
+        self.embedding_columns = embedding_columns
+        self.ner_columns = ner_columns
+        self.threshold_method = kwargs.get('threshold_method', 'idf')
+        self.tui_threshold = kwargs.get('tui_threshold', 0)
+        self.cui_threshold = kwargs.get('cui_threshold', 0)
+        self.entities_threshold = kwargs.get('entities_threshold', 0)
+
+        self.stopwatch = Stopwatch()
+
+        # self.store_similarity_matrices = kwargs.get('store_similarity_matrices', True)
+        # self.similarity_matrices_path = kwargs.get('similarity_matrices_path', True)
+
+# önce load veya split çalıştırılmalı
+    def produce_inputs(self):
+        items = []
+        y_train_label, y_test_label = stack(self.y_train), stack(self.y_test)
+        for column in self.__similarity_related_columns__:
+            train_data, test_data = stack(
+                self.X_train[column]), stack(self.X_test[column])
+            items.append([f'{column}', np.nan_to_num(train_data),
+                          y_train_label, np.nan_to_num(test_data), y_test_label])
+        for column in self.embedding_columns:
+            train_data, test_data = stack(
+                self.X_train[column+'_embedding']), stack(self.X_test[column+'_embedding'])
+            items.append([f'{column}_embedding', train_data,
+                          y_train_label, test_data, y_test_label])
+        return items
+
+##remove this function
+    def generate_sim_matrices(self, chemical_properties_df, two_d_dict):
+
+        jaccard_sim_dict = {}
+        sim_matrix_gen = SimilarityMatrixGenerator()
+
+        for column in self.__similarity_related_columns__:
+            key = '2D_'+column
+            jaccard_sim_dict[column] = sim_matrix_gen.create_jaccard_similarity_matrices(
+                two_d_dict[key])
+
+        drugbank_ids = chemical_properties_df['id'].to_list()
+
+        similarity_matrices = {}
+
+        for column in self.__similarity_related_columns__:
+            sim_matrix = jaccard_sim_dict[column]
+            jaccard_sim_feature = {}
+            for i in range(len(drugbank_ids)):
+                jaccard_sim_feature[drugbank_ids[i]] = sim_matrix[i]
+            similarity_matrices[column] = jaccard_sim_feature
+
+        return similarity_matrices
+
+    def generate_sim_matrices_new(self, chemical_properties_df):
+        self.stopwatch.reset()
+        self.stopwatch.start()
+        jaccard_sim_dict = {}
+        sim_matrix_gen = SimilarityMatrixGenerator()
+
+        for column in self.__similarity_related_columns__:
+            # key = '2D_'+column
+            key = column
+            jaccard_sim_dict[column] = sim_matrix_gen.create_jaccard_similarity_matrices(
+                self.generated_vectors[key])
+        self.stopwatch.stop()
+        print(f'similarity_matrix_generation_part_1: {self.stopwatch.elapsed}')
+
+        self.stopwatch.reset()
+        self.stopwatch.start()
+        similarity_matrices = {}
+        drugbank_ids = chemical_properties_df['id'].to_list()
+        new_columns = {}
+        for idx in range(len(drugbank_ids)):
+            new_columns[idx] = drugbank_ids[idx]
+        for column in self.__similarity_related_columns__:
+            new_df = pd.DataFrame.from_dict(jaccard_sim_dict[column])
+            new_df = new_df.rename(index=new_columns, columns=new_columns)
+            similarity_matrices[column] = new_df
+        self.stopwatch.stop()
+        print(f'similarity_matrix_generation_part_2: {self.stopwatch.elapsed}')
+        return similarity_matrices
+
+    # matris formuna çevirmek için
+    def transform_2d(self, chemical_properties_df):
+        two_d_dict = {}
+        for column in self.__similarity_related_columns__:
+            key = '2D_'+column
+            new_column = column + '_vectors'
+            two_d_dict[key] = np.stack(
+                chemical_properties_df[new_column].to_numpy())
+
+        return two_d_dict
+
+    #todo dictionary içinde ndarray dönsün
+    def generate_vectors(self, chemical_properties_df):
+        self.stopwatch.reset()
+        self.stopwatch.start()
+        vectorGenerator = VectorGenerator(chemical_properties_df)
+
+        new_columns = [
+            c+'_vectors' for c in self.__similarity_related_columns__]
+        self.generated_vectors = vectorGenerator.generate_feature_vectors(
+            self.__similarity_related_columns__)
+
+        # for column, new_column in zip(self.__similarity_related_columns__, new_columns):
+        #     chemical_properties_df.loc[:,
+        #                                new_column] = generated_vectors[column]
+        # self.generated_vectors = generated_vectors
+        self.stopwatch.stop()
+        print(f'vector_generation: {self.stopwatch.elapsed}')
+
+
+##remove this function
+    def sim(self,chemical_properties_df):
+        self.stopwatch.reset()
+        self.stopwatch.start()
+        from scipy.spatial.distance import pdist        
+        sim_matrix_gen = SimilarityMatrixGenerator()
+
+        drugbank_ids = chemical_properties_df['id'].to_list()
+        similarity_matrices = {}
+        for column in self.__similarity_related_columns__:
+            df = pd.DataFrame(np.stack(chemical_properties_df[f'{column}_vectors'].values), index = drugbank_ids)
+        #   similarity_matrices[column] = 1 - pdist(df.to_numpy(), metric='jaccard')
+            similarity_matrices[column] = sim_matrix_gen.create_jaccard_similarity_matrices(df.to_numpy())
+        self.stopwatch.stop()
+        print(f'sim: {self.stopwatch.elapsed}')
+        return similarity_matrices
+
+# import pandas as pd
+# a = [[0,0,1],[0,0,1],[0,0,0]]
+# s = pd.Series(a)
+# # print(np.vstack(s.to_numpy()))
+# l = np.argmax(np.vstack(s.to_numpy()),axis = 1)
+# l
+    def split_dataset(self,
+                      fold_size=5,
+                      shuffle=True,
+                      test_size=0.2,
+                      save_indexes=False):
+        save_path = self.index_path
+        self.prep()
+        X = self.dataframe.drop('class', axis=1)
+        y = self.dataframe['class']
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, shuffle=shuffle, test_size=test_size, stratify=np.argmax(np.vstack(y.to_numpy()),axis = 1))
+        # k_fold = KFold(n_splits=fold_size, shuffle=shuffle, random_state=1)
+        # folds = k_fold.split(X_train)
+
+        k_fold = StratifiedKFold(n_splits=fold_size, shuffle=shuffle, random_state=1)
+        folds = k_fold.split(X_train, np.argmax(np.vstack(y_train.to_numpy()),axis = 1))
+        train_idx_arr = []
+        val_idx_arr = []
+        for i, (train_index, val_index) in enumerate(folds):
+            train_idx_arr.append(train_index)
+            val_idx_arr.append(val_index)
+
+        if save_indexes:
+            # train_pairs = [row['id1'].join(',').row['id2'] for index, row in X_train.iterrows()]
+            self.__save_indexes__(
+                save_path, 'train_indexes.txt', X_train['index'].values)
+            self.__save_indexes__(
+                save_path, 'test_indexes.txt',  X_test['index'].values)
+            # self.__save_indexes__(
+            #     save_path, 'train_indexes.txt', X_train.index.values)
+            # self.__save_indexes__(
+            #     save_path, 'test_indexes.txt',  X_test.index.values)
+
+            for i, (train_idx, val_idx) in enumerate(zip(train_idx_arr, val_idx_arr)):
+                self.__save_indexes__(
+                    save_path, f'train_fold_{i}.txt', train_idx)
+                self.__save_indexes__(
+                    save_path, f'validation_fold_{i}.txt', val_idx)
+
+        self.X_train = X_train
+        self.X_test = X_test
+        self.y_train = y_train
+        self.y_test = y_test
+        self.train_indexes = X_train.index
+        self.test_indexes = X_test.index
+        self.train_idx_arr = train_idx_arr
+        self.val_idx_arr = val_idx_arr
+        return X_train, X_test, y_train, y_test, X_train.index, X_test.index, train_idx_arr, val_idx_arr
+
+    def __get_indexes__(self, path):
+        train_index_path = path+'/train_indexes.txt'
+        test_index_path = path+'/test_indexes.txt'
+        train_fold_files = f'{path}/train_fold_*.txt'
+        val_fold_files = f'{path}/validation_fold_*.txt'
+        train_idx_arr = []
+        val_idx_arr = []
+        with open(train_index_path, 'r', encoding="utf8") as f:
+            train_idx_all = [int(r) for r in f.readlines()]
+        with open(test_index_path, 'r', encoding="utf8") as f:
+            test_idx_all = [int(r) for r in f.readlines()]
+
+        for filepath in glob.glob(train_fold_files):
+            with open(filepath, 'r', encoding="utf8") as f:
+                train_idx = [int(r) for r in f.readlines()]
+                train_idx_arr.append(train_idx)
+        for filepath in glob.glob(val_fold_files):
+            with open(filepath, 'r', encoding="utf8") as f:
+                val_idx = [int(r) for r in f.readlines()]
+                val_idx_arr.append(val_idx)
+        return train_idx_all, test_idx_all, train_idx_arr, val_idx_arr
+
+    def __save_indexes__(self, path, filename, indexes):
+        create_folder_if_not_exists(path)
+        file_path = path + '/'+filename
+        str_indexes = [str(index) for index in indexes]
+        with open(file_path, 'w') as f:
+            f.write('\n'.join(str_indexes))
+
+    # @abstractmethod
+    # def prep(self):
+    #     pass
+
+    # @abstractmethod
+    # def load(self):
+    #     pass
+
+# her bir metin tipi için embedding oluşturursan burayı düzenle
+    def prep(self):
+        if self.embedding_columns:
+            zip_helper = ZipHelper()
+            zip_helper.extract(str(HERE.joinpath('zips/embeddings')),
+                               str(HERE.joinpath('zips/embeddings')))
+
+        embedding_dict = dict()
+        for embedding_column in self.embedding_columns:
+            embedding_file = HERE.joinpath(
+                f'zips/embeddings/{embedding_column}_embeddings.pkl')
+            embedding_values = pd.read_pickle(embedding_file)
+            d = embedding_values.apply(
+                lambda x: {x.id: x[f'{embedding_column}_embedding']}, axis=1)
+            x = {k: v for l in d.values.tolist() for k, v in l.items()}
+            embedding_dict[embedding_column] = x
+
+        self.ner_df = CTakesNER().load()
+        drug_names = self.drugs_df['name'].to_list()
+        drug_ids = self.drugs_df['id'].to_list()
+
+        # self.ddis_df = self.ddis_df[(self.ddis_df['name1'].isin(drug_names)) & (
+        #     self.ddis_df['name2'].isin(drug_names))]
+
+        filtered_df = self.drugs_df
+        filtered_ner_df = self.ner_df[self.ner_df['drugbank_id'].isin(
+            drug_ids)]
+        filtered_ner_df = self.ner_df.copy()
+
+        combined_df = filtered_df.copy()
+        # TODO: eğer kullanılan veri setinde tui, cui veya entity bilgileri yoksa o veri setine bu sütunları eklemek için aşağısı gerekli
+
+        idf_calc = IDF(filtered_ner_df, [f for f in filtered_ner_df.keys()])
+        idf_calc.calculate()
+        idf_scores_df = idf_calc.to_dataframe()
+
+        for key in filtered_ner_df.keys():
+            threshold = 0
+            if key.startswith('tui'):
+                threshold = self.tui_threshold
+            if key.startswith('cui'):
+                threshold = self.cui_threshold
+            if key.startswith('entities'):
+                threshold = self.entities_threshold
+            combined_df[key] = filtered_ner_df[key]
+            valid_codes = idf_scores_df[idf_scores_df[key] > threshold].index
+
+            # print(f'{key}: valid code size = {len(valid_codes)}')
+            combined_df[key] = combined_df[key].apply(lambda items:
+                                                      [item for item in items if item in valid_codes])
+
+        # Yukarıdaki koda evrildi
+        # combined_df['tui_description'] = filtered_ner_df['description_tuis']
+        # combined_df['cui_description'] = filtered_ner_df['description_cuis']
+        # combined_df['entities_description'] = filtered_ner_df['description_entities']
+
+        # tui_idf = IDF(combined_df['tui_description'], self.tui_threshold)
+        # cui_idf = IDF(combined_df['cui_description'], self.cui_threshold)
+        # entities_idf = IDF(
+        #     combined_df['entities_description'], self.entities_threshold)
+
+        # tui_idf.calculate()
+        # cui_idf.calculate()
+        # entities_idf.calculate()
+
+        # valid_tui_codes = tui_idf.find_items_over_threshold()
+        # valid_cui_codes = cui_idf.find_items_over_threshold()
+        # valid_entities_codes = entities_idf.find_items_over_threshold()
+
+        # combined_df['tui_description'] = combined_df['tui_description'].apply(lambda items:
+        #                                                                       [item for item in items if item in valid_tui_codes])
+        # combined_df['cui_description'] = combined_df['cui_description'].apply(lambda items:
+        #                                                                       [item for item in items if item in valid_cui_codes])
+        # combined_df['entities_description'] = combined_df['entities_description'].apply(lambda items:
+        #                                                                                 [item for item in items if item in valid_entities_codes])
+
+        moved_columns = ['id']
+        moved_columns.extend(self.__similarity_related_columns__)
+        chemical_properties_df = combined_df[moved_columns]
+
+        chemical_properties_df = chemical_properties_df.fillna("").apply(list)
+
+        # generate vectors dictionary içinde ndarray dönecek
+        self.generate_vectors(chemical_properties_df)
+
+        # two_d_dict = self.transform_2d(chemical_properties_df)
+
+        similarity_matrices = self.generate_sim_matrices_new(
+            chemical_properties_df)
+
+        # similarity_matrices = self.sim(chemical_properties_df)
+
+        event_categories = self.ddis_df['event_category']
+        labels = event_categories.tolist()
+        lb = LabelBinarizer()
+        lb.fit(labels)
+        classes = lb.transform(labels)
+
+        # def similarity_lambda_fnc(row, value):
+        #     if row['id1'] in value and row['id2'] in value:
+        #         return value[row['id1']][row['id2']]
+            
+        def similarity_lambda_fnc(row, value):
+            if row['id1'] in value:
+                return value[row['id1']]
+
+        def lambda_fnc(row, value):
+            if row['id1'] in value and row['id2'] in value:
+                return np.float16(np.hstack(
+                    (value[row['id1']], value[row['id2']])))
+                # return np.hstack(
+                #     (value[row['id1']], value[row['id2']]), dtype=np.float16)
+
+        def x_fnc(row, embedding_values, embedding_column):
+            # first = embedding_values[embedding_values.id == row['id1']]
+            # second = embedding_values[embedding_values.id == row['id2']]
+            # v1 = first.iloc[0][embedding_column+'_embedding']
+            # v2 = second.iloc[0][embedding_column+'_embedding']
+            v1 = embedding_dict[embedding_column][row['id1']]
+            v2 = embedding_dict[embedding_column][row['id2']]
+            # v1 = embedding_dict[row['id1']][embedding_column+'_embedding']
+            # v2 = embedding_dict[row['id2']][embedding_column+'_embedding']
+            return np.float16(np.hstack(
+                (v1, v2)))
+
+        for key, value in similarity_matrices.items():
+
+            print(f'sim matrix: {key}')
+            self.ddis_df[key] = self.ddis_df.apply(
+                lambda_fnc, args=(value,), axis=1)
+            print(self.ddis_df[key].head())
+
+        for embedding_column in self.embedding_columns:
+            print(f"concat {embedding_column} embeddings")
+            # column_embeddings_dict = embedding_values[embedding_column]
+            self.ddis_df[embedding_column+'_embedding'] = self.ddis_df.apply(
+                x_fnc, args=(embedding_values, embedding_column), axis=1)
+
+        self.dataframe = self.ddis_df.copy()
+        self.dataframe['class'] = list(classes)
+        print(self.dataframe.shape)
+
+    def load(self):
+        if self.index_path == None:
+            raise Exception(
+                "There is no index path, please call split function")
+
+        # prep - split - load
+        train_idx_all, test_idx_all, train_idx_arr, val_idx_arr = self.__get_indexes__(
+            self.index_path)
+
+        self.prep()
+        train = self.dataframe[self.dataframe['index'].isin(train_idx_all)]
+        test = self.dataframe[self.dataframe['index'].isin(test_idx_all)]
+
+        self.X_train = train.drop('class', axis=1)
+        self.y_train = train['class']
+        self.X_test = test.drop('class', axis=1)
+        self.y_test = test['class']
+
+        self.train_indexes = self.X_train.index
+        self.test_indexes = self.X_test.index
+        self.train_idx_arr = train_idx_arr
+        self.val_idx_arr = val_idx_arr
+
+        return self.X_train, self.X_test, self.y_train, self.y_test, self.X_train.index, self.X_test.index, train_idx_arr, val_idx_arr

ddi_fw/datasets/db_utils.py

@@ -0,0 +1,204 @@
+
+from sqlite3 import Error
+import sqlite3
+import pandas as pd
+import numpy as np
+
+
+def create_connection(db_file=r"./event.db"):
+    """ create a database connection to the SQLite database
+        specified by db_file
+    :param db_file: database file
+    :return: Connection object or None
+    """
+    conn = None
+    try:
+        conn = sqlite3.connect(db_file)
+    except Error as e:
+        print(e)
+
+    return conn
+
+
+def select_all_drugs(conn):
+    cur = conn.cursor()
+    cur.execute(
+        '''select "index", id, name, target, enzyme, pathway, smile from drug''')
+    rows = cur.fetchall()
+    return rows
+
+
+def select_all_drugs_as_dataframe(conn):
+    headers = ['index','id', 'name', 'target', 'enzyme', 'pathway', 'smile']
+    rows = select_all_drugs(conn)
+    df = pd.DataFrame(columns=headers, data=rows)
+    df['enzyme'] = df['enzyme'].apply(lambda x: x.split('|'))
+    df['target'] = df['target'].apply(lambda x: x.split('|'))
+    df['pathway'] = df['pathway'].apply(lambda x: x.split('|'))
+    df['smile'] = df['smile'].apply(lambda x: x.split('|'))
+    return df
+
+
+def select_all_events(conn):
+    """
+    Query all rows in the event table
+    :param conn: the Connection object
+    :return:
+    """
+    cur = conn.cursor()
+    cur.execute("select * from event")
+
+    rows = cur.fetchall()
+    return rows
+
+
+def select_all_events_as_dataframe(conn):
+    headers = ["index", "id1", "name1", "id2", "name2", "event_category"]
+    rows = select_all_events(conn)
+    return pd.DataFrame(columns=headers, data=rows)
+
+
+def select_events_with_category(conn):
+    sql = '''select id1, name1, id2, name2, mechanism || ' ' ||action from event ev
+            join extraction ex
+            on ev.name1 = ex.drugA and ev.name2 = ex.drugB
+            union 
+            select id1, name1, id2, name2, mechanism || ' ' ||action from event ev
+            join extraction ex
+            on ev.name1 = ex.drugB and ev.name2 = ex.drugA
+            '''
+    cur = conn.cursor()
+    cur.execute(sql)
+
+    rows = cur.fetchall()
+
+    headers = ['id1', 'name1', 'id2', 'name2', 'event_category']
+    return pd.DataFrame(columns=headers, data=rows)
+
+
+def select_all_interactions_tuple_as_dataframe(conn):
+    cur = conn.cursor()
+    cur.execute("select id1, id2 from event")
+    rows = cur.fetchall()
+    headers = ['id1', 'id2']
+
+    return pd.DataFrame(columns=headers, data=rows)
+
+
+def select_ddi_pairs(conn):
+    cur = conn.cursor()
+    cur.execute('''
+        select d1.[index] as Drug1Index, d2.[index] as Drug2Index, 1 from event e 
+        join drug d1 on e.id1 = d1.id 
+        join drug d2 on e.id2 = d2.id 
+    ''')
+    rows = cur.fetchall()
+    return rows
+
+
+def select_ddi_pairs_as_dataframe(conn):
+    headers = ["Drug1Index", "Drug2Index", "Interaction"]
+    rows = select_ddi_pairs(conn)
+    return pd.DataFrame(columns=headers, data=rows)
+
+
+def get_interactions(conn):
+    cur = conn.cursor()
+    cur.execute('''
+        select
+            drug_1_id,
+            drug_1,
+            drug_2_id,
+            drug_2,
+            mechanism_action, 
+            interaction, 
+            masked_interaction
+        from _Interactions
+                ''')
+
+    rows = cur.fetchall()
+
+    headers = ['id1', 'name1', 'id2', 'name2',
+               'event_category', 'interaction', 'masked_interaction']
+    df = pd.DataFrame(columns=headers, data=rows)
+    return df
+
+
+def get_extended_version(conn):
+    cur = conn.cursor()
+    cur.execute('''
+        select 
+        _Drugs."index",
+        drugbank_id,
+        _Drugs.name,
+        description,
+        synthesis_reference,
+        indication,
+        pharmacodynamics,
+        mechanism_of_action,
+        toxicity,
+        metabolism,
+        absorption,
+        half_life,
+        protein_binding,
+        route_of_elimination,
+        volume_of_distribution,
+        clearance,
+        smiles,
+        smiles_morgan_fingerprint,
+        enzymes_polypeptides,
+        targets_polypeptides
+                
+        from drug 
+        join _Drugs on drug.id = _Drugs.drugbank_id  
+        where 
+                targets_polypeptides is not null and 
+                enzymes_polypeptides is not null and 
+                smiles_morgan_fingerprint is not null
+''')
+    # pathway is absent
+
+    rows = cur.fetchall()
+    headers = ['index', 'id', 'name', 'description', 'synthesis_reference', 'indication', 'pharmacodynamics', 'mechanism_of_action', 'toxicity', 'metabolism', 'absorption', 'half_life',
+               'protein_binding', 'route_of_elimination', 'volume_of_distribution', 'clearance', 'smiles_notation', 'smile', 'enzyme', 'target']
+    df = pd.DataFrame(columns=headers, data=rows)
+    df['smile'] = df['smile'].apply(lambda x:
+                                    np.fromstring(
+                                        x.replace(
+                                            '\n', '')
+                                        .replace('[', '')
+                                        .replace(']', '')
+                                        .replace('  ', ' '), sep=','))
+    df['enzyme'] = df['enzyme'].apply(
+        lambda x: x.split('|'))
+    df['target'] = df['target'].apply(
+        lambda x: x.split('|'))
+    return df
+
+
+# SELECT
+#  CASE
+#  WHEN masked_interaction like '%'+drug_1+'%' THEN drug_1
+#  WHEN masked_interaction like '%'+drug_2+'%' THEN drug_2
+#  Else drug_2
+#  END AS Absent,
+
+# drug_1, drug_2,
+# masked_interaction
+
+# from _Interactions
+# where LENGTH(masked_interaction) = LENGTH(REPLACE(masked_interaction, 'DRUG', ''))
+# or LENGTH(masked_interaction) = LENGTH(REPLACE(masked_interaction, 'DRUG', '')) + 4
+
+if __name__ == "__main__":
+    conn = create_connection(r"./event-extended.db")
+    extended_version_df = get_extended_version(conn)
+
+    df = select_all_events_as_dataframe(conn)
+    print(df.head())
+
+    events_with_category_df = select_events_with_category(conn)
+    print(events_with_category_df.head())
+
+    u = events_with_category_df['event_category'].unique()
+    print(len(u))