ddi-fw 0.0.78__py3-none-any.whl → 0.0.79__py3-none-any.whl

ddi_fw/experiments/__init__.py

@@ -1,5 +1,4 @@
 from .tensorflow_helper import TFMultiModal, TFSingleModal,Result
 from .evaluation_helper import evaluate, Metrics
 from .pipeline import Experiment
-from .pipeline_ner import NerParameterSearch
-from .ml_helper import ModelWrapper,MultiModalRunner
+from .pipeline_ner import NerParameterSearch

ddi_fw/ml/__init__.py

@@ -0,0 +1,4 @@
+from .ml_helper import MultiModalRunner
+from .model_wrapper import ModelWrapper,Result
+from .tensorflow_wrapper import TFModelWrapper
+from .pytorch_wrapper import PTModelWrapper

ddi_fw/{experiments → ml}/ml_helper.py

@@ -1,6 +1,7 @@
 from typing import Dict, List, Tuple
 from matplotlib import pyplot as plt
-from ddi_fw.experiments.ml_pt import PTModelWrapper
+from ddi_fw.ml.model_wrapper import Result
+from ddi_fw.ml.pytorch_wrapper import PTModelWrapper
 from ddi_fw.experiments.ml_tf import TFModelWrapper
 import tensorflow as tf
 from tensorflow import keras
@@ -28,37 +29,7 @@ np.random.seed(2)
 np.set_printoptions(precision=4)
 
 
-class Result:
-    def __init__(self) -> None:
-        self.log_dict = {}
-        self.metric_dict = {}
 
-    def add_log(self, key, logs):
-        self.log_dict[key] = logs
-
-    def add_metric(self, key, metrics):
-        self.metric_dict[key] = metrics
-
-
-class ModelWrapper:
-    def __init__(self, date, descriptor, model_func, batch_size=128, epochs=100):
-        self.date = date
-        self.descriptor = descriptor
-        self.model_func = model_func
-        self.batch_size = batch_size
-        self.epochs = epochs
-
-    def set_data(self, train_idx_arr, val_idx_arr, train_data, train_label, test_data, test_label):
-        self.train_idx_arr = train_idx_arr
-        self.val_idx_arr = val_idx_arr
-        self.train_data = train_data
-        self.train_label = train_label
-        self.test_data = test_data
-        self.test_label = test_label
-# https://github.com/mlflow/mlflow/blob/master/examples/tensorflow/train.py
-
-    def predict(self) -> Tuple[Dict[str, float], Metrics, List[float]]:
-        pass
  
 
 class MultiModalRunner:
@@ -101,7 +72,7 @@ class MultiModalRunner:
             self.level_0_run_id = run.info.run_id
             for item in self.items:
                 print(item[0])
-                T =self.__create_multi_modal(self.library)
+                T = self.__create_multi_modal(self.library)
                 single_modal=T(self.date, item[0], self.model_func, self.batch_size, self.epochs)
                 single_modal.set_data(
                     self.train_idx_arr, self.val_idx_arr, item[1], item[2], item[3], item[4])

ddi_fw/ml/model_wrapper.py

@@ -0,0 +1,35 @@
+from typing import Dict, List, Tuple
+
+from ddi_fw.experiments.evaluation_helper import Metrics
+
+class Result:
+    def __init__(self) -> None:
+        self.log_dict = {}
+        self.metric_dict = {}
+
+    def add_log(self, key, logs):
+        self.log_dict[key] = logs
+
+    def add_metric(self, key, metrics):
+        self.metric_dict[key] = metrics
+
+
+class ModelWrapper:
+    def __init__(self, date, descriptor, model_func, batch_size=128, epochs=100):
+        self.date = date
+        self.descriptor = descriptor
+        self.model_func = model_func
+        self.batch_size = batch_size
+        self.epochs = epochs
+
+    def set_data(self, train_idx_arr, val_idx_arr, train_data, train_label, test_data, test_label):
+        self.train_idx_arr = train_idx_arr
+        self.val_idx_arr = val_idx_arr
+        self.train_data = train_data
+        self.train_label = train_label
+        self.test_data = test_data
+        self.test_label = test_label
+# https://github.com/mlflow/mlflow/blob/master/examples/tensorflow/train.py
+
+    def predict(self) -> Tuple[Dict[str, float], Metrics, List[float]]:
+        pass

ddi_fw/{experiments/ml_pt.py → ml/pytorch_wrapper.py}

@@ -1,7 +1,7 @@
 import mlflow
 import torch
-from ddi_fw.experiments.ml_helper import ModelWrapper
 from ddi_fw.experiments.evaluation_helper import evaluate
+from ddi_fw.ml.model_wrapper import ModelWrapper
 
 
 class PTModelWrapper(ModelWrapper):

ddi_fw/{experiments/ml_tf.py → ml/tensorflow_wrapper.py}

@@ -1,6 +1,5 @@
-from typing import Dict, List, Tuple
 from matplotlib import pyplot as plt
-from ddi_fw.experiments.ml_helper import ModelWrapper
+from ddi_fw.ml.model_wrapper import ModelWrapper
 import tensorflow as tf
 from tensorflow import keras
 from keras.models import Model, Sequential

ddi_fw/pipeline/__init__.py

@@ -0,0 +1,3 @@
+from .pipeline import Pipeline
+from .multi_pipeline import MultiPipeline
+from multi_modal_combination_strategy import CombinationStrategy,CustomCombinationStrategy

ddi_fw/pipeline/multi_modal_combination_strategy.py

@@ -0,0 +1,39 @@
+import itertools
+
+
+class CombinationStrategy():
+    def generate(self):
+        pass
+
+
+class CustomCombinationStrategy(CombinationStrategy):
+    def __init__(self, **kwargs_combination_params):
+        # kwargs fonksiyona da alınabilir
+        self.group1 = kwargs_combination_params.get("group_1", None)
+        self.group2 = kwargs_combination_params.get("group_2", None)
+
+    def generate(self):
+        # Handle edge cases
+        if not self.group_1 or not self.group_2:
+            raise ValueError(
+                f"Parameters of combination strategy could not be empty.")
+            # return []  # Return an empty list if either group is empty
+        # combinations = []
+        # for j in self.group2:
+        #     extended_item_group_1 = self.group_1.copy()
+        #     extended_item_group_1.append(j)
+        #     for i in range(2, len(extended_item_group_1) + 1):
+        #         combinations.extend(list(itertools.combinations(extended_item_group_1, i))) #all
+        # combinations = list(set(combinations))
+
+        combinations = set()  # Use a set to avoid duplicates directly
+        for j in self.group_2:
+            extended_item_group_1 = self.group_1.copy()
+            extended_item_group_1.append(j)
+            # Generate combinations of all lengths from 2 to len(group_1 + 1)
+            for i in range(2, len(extended_item_group_1) + 1):
+                combinations.update(itertools.combinations(
+                    extended_item_group_1, i))  # Add combinations
+
+        # Convert set back to list (if needed) and return
+        return list(combinations)

ddi_fw/pipeline/multi_pipeline.py

@@ -0,0 +1,111 @@
+import json
+from pipeline import Pipeline
+import importlib
+
+
+def load_config(file_path):
+    with open(file_path, 'r') as file:
+        config = json.load(file)
+    return config
+
+
+def get_import(full_path_of_import):
+    """Dynamically imports an object from a module given its full path.
+
+    Args:
+        full_path_of_import (str): The full path of the import (e.g., 'module.submodule.ClassName').
+
+    Returns:
+        object: The imported object.
+
+    Raises:
+        ImportError: If the module cannot be imported.
+        AttributeError: If the attribute does not exist in the module.
+    """
+    if not full_path_of_import:
+        raise ValueError("The import path cannot be empty.")
+
+    parts = full_path_of_import.split('.')
+    import_name = parts[-1]
+    module_name = ".".join(parts[:-1]) if len(parts) > 1 else ""
+
+    try:
+        module = importlib.import_module(module_name)
+        return getattr(module, import_name)
+    except ModuleNotFoundError as e:
+        raise ImportError(f"Module '{module_name}' could not be found.") from e
+    except AttributeError as e:
+        raise AttributeError(
+            f"'{module_name}' has no attribute '{import_name}'") from e
+
+
+class MultiPipeline():
+    def __init__(self, experiments_config_file):
+        self.experiments_config = load_config(experiments_config_file)
+        self.items = []
+
+    def __create_pipeline(self, config):
+        library = config["library"]
+        batch_size = config["batch_size"]
+        epochs = config["epochs"]
+
+        # dataset_module = config["dataset_module"]
+        # dataset_name = config["dataset_name"]
+
+        experiment_name = config["experiment_name"]
+        experiment_description = config["experiment_description"]
+        experiment_tags = config["experiment_tags"]
+        tracking_uri = config["tracking_uri"]
+        artifact_location = config["artifact_location"]
+        columns = config["columns"]
+        ner_data_file = config["ner_data_file"]
+        ner_threshold = config["ner_threshold"]
+        vector_db_persist_directory = config["vector_db_persist_directory"]
+        vector_db_collection_name = config["vector_db_collection_name"]
+        embedding_pooling_strategy = get_import(
+            config["embedding_pooling_strategy_type"])
+        # Dynamically import the model and dataset classes
+        model_type = get_import(config["model_type"])
+        dataset_type = get_import(config["dataset_type"])
+        combination_type = get_import(config["combination_strategy"]["type"])
+        kwargs_combination_params = config["combination_strategy"]["params"]
+
+        # # Instantiate the classes
+        # model_instance = model_class()
+        # dataset_instance = dataset_class()
+        return {
+            "name": experiment_name,
+            "library": library,
+            "batch_size": batch_size,
+            "epochs": epochs,
+            "model_type": model_type,
+            "pipeline": Pipeline(
+                library=library,
+                experiment_name=experiment_name,
+                experiment_description=experiment_description,
+                experiment_tags=experiment_tags,
+                artifact_location=artifact_location,
+                tracking_uri=tracking_uri,
+                dataset_type=dataset_type,
+                columns=columns,
+                vector_db_persist_directory=vector_db_persist_directory,
+                vector_db_collection_name=vector_db_collection_name,
+                embedding_pooling_strategy_type=embedding_pooling_strategy,
+                ner_data_file=ner_data_file,
+                ner_threshold=ner_threshold,
+                combinations=combination_type(**kwargs_combination_params).generate())}
+
+    def build(self):
+        for config in self.experiments_config['experiments']:
+            item = self.__create_pipeline(config)
+            self.items.append(item)
+
+    def run(self):
+        for item in self.items:
+            print(f"{item['name']} is running")
+            pipeline = item['pipeline']
+            model_type = item['model_type']
+            batch_size = item['batch_size']
+            epochs = item['epochs']
+            pipeline.build()
+            pipeline.run(model_type, epochs=epochs, batch_size=batch_size)

ddi_fw/pipeline/pipeline.py

@@ -0,0 +1,126 @@
+import numpy as np
+import pandas as pd
+import chromadb
+from collections import defaultdict
+from ddi_fw.ner.ner import CTakesNER
+from ddi_fw.langchain.embeddings import PoolingStrategy
+from ddi_fw.datasets import BaseDataset, DDIMDLDataset
+from ddi_fw.langchain.embeddings import SumPoolingStrategy
+import mlflow
+from ml import MultiModalRunner
+
+
+class Pipeline:
+    def __init__(self,
+                 library='TF',
+                 experiment_name=None,
+                 experiment_description=None,
+                 experiment_tags=None,
+                 artifact_location=None,
+                 tracking_uri=None,
+                 dataset_type: BaseDataset = None,
+                 columns=None,
+                 embedding_dict=None,
+                 vector_db_persist_directory=None,
+                 vector_db_collection_name=None,
+                 embedding_pooling_strategy_type: PoolingStrategy = None,
+                 ner_data_file=None,
+                 ner_threshold=None,
+                 combinations=None,
+                 model=None):
+        self.library = library
+        self.experiment_name = experiment_name
+        self.experiment_description = experiment_description
+        self.experiment_tags = experiment_tags
+        self.artifact_location = artifact_location
+        self.tracking_uri = tracking_uri
+        self.dataset_type = dataset_type
+        self.columns = columns
+        self.embedding_dict = embedding_dict
+        self.vector_db_persist_directory = vector_db_persist_directory
+        self.vector_db_collection_name = vector_db_collection_name
+        self.embedding_pooling_strategy_type = embedding_pooling_strategy_type
+        self.ner_data_file = ner_data_file
+        self.ner_threshold = ner_threshold
+        self.combinations = combinations
+        self.model = model
+
+    def build(self):
+        # 'enzyme','target','pathway','smile','all_text','indication', 'description','mechanism_of_action','pharmacodynamics', 'tui', 'cui', 'entities'
+        kwargs = {"columns": self.columns}
+        for k, v in self.ner_threshold.items():
+            kwargs[k] = v
+        if self.embedding_dict == None:
+            if self.vector_db_persist_directory:
+                self.vector_db = chromadb.PersistentClient(
+                    path=self.vector_db_persist_directory)
+                self.collection = self.vector_db.get_collection(
+                    self.vector_db_collection_name)
+                dictionary = self.collection.get(
+                    include=['embeddings', 'metadatas'])
+
+                embedding_dict = defaultdict(lambda: defaultdict(list))
+
+                for metadata, embedding in zip(dictionary['metadatas'], dictionary['embeddings']):
+                    embedding_dict[metadata["type"]
+                                   ][metadata["id"]].append(embedding)
+
+                embedding_size = dictionary['embeddings'].shape[1]
+        else:
+            embedding_dict = self.embedding_dict
+            embedding_size = list(embedding_dict['all_text'].values())[
+                0][0].shape
+
+        pooling_strategy = self.embedding_pooling_strategy_type()
+
+        self.ner_df = CTakesNER().load(
+            filename=self.ner_data_file) if self.ner_data_file else None
+
+        self.dataset = self.dataset_type(
+            embedding_dict=embedding_dict,
+            embedding_size=embedding_size,
+            embeddings_pooling_strategy=pooling_strategy,
+            ner_df=self.ner_df, **kwargs)
+
+        X_train, X_test, y_train, y_test, X_train.index, X_test.index, train_idx_arr, val_idx_arr = self.dataset.load()
+
+        self.dataframe = self.dataset.dataframe
+        # dataframe.dropna()
+        self.X_train = self.dataset.X_train
+        self.X_test = self.dataset.X_test
+        self.y_train = self.dataset.y_train
+        self.y_test = self.dataset.y_test
+        self.train_idx_arr = self.dataset.train_idx_arr
+        self.val_idx_arr = self.dataset.val_idx_arr
+        # Logic to set up the experiment
+        self.items = self.dataset.produce_inputs()
+
+        unique_classes = pd.unique(self.dataframe['event_category'])
+        event_num = len(unique_classes)
+        # droprate = 0.3
+        vector_size = self.dataset.drugs_df.shape[0]
+
+        print("Building the experiment with the following settings:")
+        print(
+            f"Name: {self.experiment_name}, Dataset: {self.dataset}, Model: {self.model}")
+        # Implement additional build logic as needed
+        return self
+
+    def run(self, model_func, batch_size=128, epochs=100):
+        mlflow.set_tracking_uri(self.tracking_uri)
+
+        if mlflow.get_experiment_by_name(self.experiment_name) == None:
+            mlflow.create_experiment(
+                self.experiment_name, self.artifact_location)
+            mlflow.set_experiment_tags(self.experiment_tags)
+        mlflow.set_experiment(self.experiment_name)
+
+        y_test_label = self.items[0][4]
+        multi_modal_runner = MultiModalRunner(
+            library=self.library, model_func=model_func, batch_size=batch_size,  epochs=epochs)
+        # multi_modal = TFMultiModal(
+        #     model_func=model_func, batch_size=batch_size,  epochs=epochs)  # 100
+        multi_modal_runner.set_data(
+            self.items, self.train_idx_arr, self.val_idx_arr, y_test_label)
+        result = multi_modal_runner.predict(self.combinations)
+        return result

{ddi_fw-0.0.78.dist-info → ddi_fw-0.0.79.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ddi_fw
-Version: 0.0.78
+Version: 0.0.79
 Summary: Do not use :)
 Author-email: Kıvanç Bayraktar <bayraktarkivanc@gmail.com>
 Maintainer-email: Kıvanç Bayraktar <bayraktarkivanc@gmail.com>

{ddi_fw-0.0.78.dist-info → ddi_fw-0.0.79.dist-info}/RECORD

@@ -55,14 +55,10 @@ ddi_fw/drugbank/drugbank_parser.py,sha256=lxUuhB0s8ef_aPNDs0V8ClKF7-KIWugNIV9gVs
 ddi_fw/drugbank/drugbank_processor.py,sha256=vmkt68n9nFLevufgGyXhOSDtTo4G1XzwT9PVncGTXtk,18127
 ddi_fw/drugbank/drugbank_processor_org.py,sha256=eO5Yset50P91qkic79RUXPoEuxRxQKFkKW0l4G29Mas,13322
 ddi_fw/drugbank/event_extractor.py,sha256=6odoZohhK7OdLF-LF0l-5BFq0_NMG_5jrFJbHrBXsI8,4600
-ddi_fw/experiments/__init__.py,sha256=QVWBpJ9x45GJmW5vM5YwpkASr5Dam88-P5DwzcdSsLs,249
+ddi_fw/experiments/__init__.py,sha256=5L2xSolpFycNnflqOMdvJSiqRB16ExA5bbVGORKFX04,195
 ddi_fw/experiments/custom_torch_model.py,sha256=iQ_R_EApzD2JCcASN8cie6D21oh7VCxaOQ45_dkiGwc,2576
 ddi_fw/experiments/evaluation_helper.py,sha256=o4-w5Xa3t4olLW4ymx_8L-Buhe5wfQEmT2bh4Zz544c,13066
-ddi_fw/experiments/ml_helper.py,sha256=1mH6IFhOG4eZ-GueTa_-8V9OATapWp_VclnCCpVwnnE,5473
-ddi_fw/experiments/ml_pt.py,sha256=VHKegdX5-RyUpNN_l6XxMc2ZrSQg4h8uuQALInQXkRg,3730
-ddi_fw/experiments/ml_tf.py,sha256=flOsVxCrok5zIlT4OHbk3NzEhtgyybMAQOENI6Itn9I,5791
 ddi_fw/experiments/pipeline.py,sha256=4ltPCcfLZ1fFpiOd8ahPognI6NLmRLzJvUqyFpn3z18,5693
-ddi_fw/experiments/pipeline_builder_pattern.py,sha256=w6x7ietk4vONCAvUfssPycaRUQIYUJsbCNNj3BTASBI,5454
 ddi_fw/experiments/pipeline_ner.py,sha256=unxEJCYrG6wEZjLmqvGdLRTMOBwELbGKkdygSpAR3b8,5043
 ddi_fw/experiments/tensorflow_helper.py,sha256=m3Mppl-tbccTMAKLpZg2YC0xpcukkyQihPw_uwAlRRY,11857
 ddi_fw/experiments/test.py,sha256=z1TfBpK75zGKpp2ZU8f6APjZlgBFthaCBN61YB9ma4o,2049
@@ -70,9 +66,18 @@ ddi_fw/langchain/__init__.py,sha256=8dBPZivc01WWaCH8sZ_UV8-XPyo74e9Qy6-fYgAiNLE,
 ddi_fw/langchain/embeddings.py,sha256=8J_SfO9pyET2W-Ltzq0_r9EchFzBsYdUabiOMma42Us,7515
 ddi_fw/langchain/sentence_splitter.py,sha256=h_bYElx4Ud1mwDNJfL7mUwvgadwKX3GKlSzu5L2PXzg,280
 ddi_fw/langchain/storage.py,sha256=uy5clVB07So2eFbRGdAKzHIPdfEk4se33cPktis7Aa4,2716
+ddi_fw/ml/__init__.py,sha256=0YubqmEpJKp3OfqlLKkD5N9L6WDWew3QEtnbdY3mqKg,180
+ddi_fw/ml/ml_helper.py,sha256=juDcTi8IEQk2D4mkY4qVX75rRM0FmksULRQzyNHKw2A,4475
+ddi_fw/ml/model_wrapper.py,sha256=ZExnsLMjHKL3BaI4aKkbyWTp8vbswLeF2_T3cZ73YpQ,1144
+ddi_fw/ml/pytorch_wrapper.py,sha256=YdwzR5qAHFNajYB_elFqDhVKRLeajaRpopNzyQ6gIIA,3725
+ddi_fw/ml/tensorflow_wrapper.py,sha256=pSeiJDuaLf9MhZVlLuLJBA-LH-H-Dl2TyYbB39iGsto,5748
 ddi_fw/ner/__init__.py,sha256=JwhGXrepomxPSsGsg2b_xPRC72AjvxOIn2CW5Mvscn0,26
 ddi_fw/ner/mmlrestclient.py,sha256=NZta7m2Qm6I_qtVguMZhqtAUjVBmmXn0-TMnsNp0jpg,6859
 ddi_fw/ner/ner.py,sha256=BEs9AFljAxOQrC2BEP1raSzRoypcfELS5UTdl4bjTqw,15863
+ddi_fw/pipeline/__init__.py,sha256=qryVi8bTsbpbMsseOuSEi1Siign0LkbFLPWiIR7OGHE,165
+ddi_fw/pipeline/multi_modal_combination_strategy.py,sha256=YkPixHVo9-4SPkY8VaWvBe1aaI5IiV4oZT4kBrm2WHQ,1635
+ddi_fw/pipeline/multi_pipeline.py,sha256=UgTEcT2UfkRKR3Ri_Nrtz9GrQNQHGOSUrw9h5AwFUMI,4356
+ddi_fw/pipeline/pipeline.py,sha256=WJnz5zEIa-9n4qEs8-1ubqTMPsLYjFGdFYJPbn92i98,5512
 ddi_fw/test/basic_test.py,sha256=fEOGcZm1ObnsDvMiXNmdmz6YCeUrGc8V0DwlSwGhsq8,376
 ddi_fw/test/combination_test.py,sha256=TWNE8sf-DSh1Q9-yRaRBc774Sn1kSMGXLwQhd2_Qynk,324
 ddi_fw/test/compress_json_test.py,sha256=BGny56YqiG-pzhMoDzLKQBQI1E7o3jU0S7VYWtclAx4,1045
@@ -89,7 +94,7 @@ ddi_fw/utils/enums.py,sha256=19eJ3fX5eRK_xPvkYcukmug144jXPH4X9zQqtsFBj5A,671
 ddi_fw/utils/py7zr_helper.py,sha256=gOqaFIyJvTjUM-btO2x9AQ69jZOS8PoKN0wetYIckJw,4747
 ddi_fw/utils/utils.py,sha256=szwnxMTDRrZoeNRyDuf3aCbtzriwtaRk4mHSH3asLdA,4301
 ddi_fw/utils/zip_helper.py,sha256=YRZA4tKZVBJwGQM0_WK6L-y5MoqkKoC-nXuuHK6CU9I,5567
-ddi_fw-0.0.78.dist-info/METADATA,sha256=2nUPcQaInXGSs6zh6fqhDHkvyyeql-B0pdXGsfAoE6Y,1966
-ddi_fw-0.0.78.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-ddi_fw-0.0.78.dist-info/top_level.txt,sha256=PMwHICFZTZtcpzQNPV4UQnfNXYIeLR_Ste-Wfc1h810,7
-ddi_fw-0.0.78.dist-info/RECORD,,
+ddi_fw-0.0.79.dist-info/METADATA,sha256=Acf-Yb4NTk6aKueaQ3tbti_Ykxm6mkzJP-mreWb8UWI,1966
+ddi_fw-0.0.79.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+ddi_fw-0.0.79.dist-info/top_level.txt,sha256=PMwHICFZTZtcpzQNPV4UQnfNXYIeLR_Ste-Wfc1h810,7
+ddi_fw-0.0.79.dist-info/RECORD,,

ddi_fw/experiments/pipeline_builder_pattern.py

@@ -1,152 +0,0 @@
-import sqlite3
-from sklearn.model_selection import train_test_split, KFold, StratifiedKFold
-from keras.models import Model, Sequential
-from keras.callbacks import EarlyStopping
-from keras.layers import Dense, Dropout, Input, Activation, BatchNormalization
-from tensorflow import keras
-from ddi_fw.experiments import TFSingleModal, TFMultiModal
-from ddi_fw.experiments import evaluate
-from sklearn.preprocessing import LabelBinarizer
-import numpy as np
-import pandas as pd
-from ddi_fw.utils import ZipHelper, Py7ZipHelper
-import os
-import chromadb
-from collections import defaultdict
-from langchain_community.vectorstores import Chroma
-from ddi_fw.ner.ner import CTakesNER
-from ddi_fw.langchain.embeddings import PoolingStrategy
-
-from ddi_fw.datasets import BaseDataset, DDIMDLDataset
-
-from ddi_fw.datasets import SumPoolingStrategy
-from keras import metrics
-from ddi_fw.experiments.evaluation_helper import evaluate
-
-import mlflow
-
-
-class Experiment:
-    def __init__(self):
-        pass
-
-    @staticmethod
-    def create():
-        return Experiment()
-
-    def name(self, name):
-        self.experiment_name = name
-        return self
-
-    def description(self, description):
-        self.experiment_description = description
-        return self
-
-    def tags(self, tags):
-        self.experiment_tags = tags
-        return self
-
-    def tracking_uri(self, uri):
-        self.tracking_uri = uri
-        return self
-
-    def dataset(self, dataset_type: BaseDataset):
-        self.dataset_type = dataset_type
-        return self
-
-    def columns(self, cols):
-        self.columns = cols
-        return self
-
-    def vectordb_collection(self, persist_directory, collection_name):
-        self.vector_db_persist_directory = persist_directory
-        self.vector_db_collection_name = collection_name
-        return self
-
-    def embedding_pooling_strategy(self, strategy_type: PoolingStrategy):
-        self.embedding_pooling_strategy_type = strategy_type
-        return self
-
-    def ner_data_file(self, ner_data_file):
-        self.ner_data_file = ner_data_file
-        self.ner_df = CTakesNER().load(filename=ner_data_file)
-        return self
-
-    def ner_threshold(self, threshold):
-        self.ner_threshold = threshold
-        return self
-
-    def combinations(self, combs):
-        self.combinations = combs
-        return self
-
-    def model(self, model):
-        self.model = model
-        return self
-
-    def build(self):
-        # 'enzyme','target','pathway','smile','all_text','indication', 'description','mechanism_of_action','pharmacodynamics', 'tui', 'cui', 'entities'
-        kwargs = {"columns": self.columns}
-        for k, v in self.ner_threshold.items():
-            kwargs[k] = v
-
-        self.vector_db = chromadb.PersistentClient(
-            path=self.vector_db_persist_directory)
-        self.collection = self.vector_db.get_collection(
-            self.vector_db_collection_name)
-        dictionary = self.collection.get(include=['embeddings', 'metadatas'])
-
-        embedding_dict = defaultdict(lambda: defaultdict(list))
-
-        for metadata, embedding in zip(dictionary['metadatas'], dictionary['embeddings']):
-            embedding_dict[metadata["type"]][metadata["id"]].append(embedding)
-
-        embedding_size = dictionary['embeddings'].shape[1]
-
-        pooling_strategy = self.embedding_pooling_strategy_type()
-
-        self.dataset = self.dataset_type(
-            embedding_dict=embedding_dict,
-            embedding_size=embedding_size,
-            embeddings_pooling_strategy=pooling_strategy,
-            ner_df=self.ner_df, kwargs=kwargs)
-
-        X_train, X_test, y_train, y_test, X_train.index, X_test.index, train_idx_arr, val_idx_arr = self.dataset.load()
-
-        self.dataframe = self.dataset.dataframe
-        # dataframe.dropna()
-        self.X_train = self.dataset.X_train
-        self.X_test = self.dataset.X_test
-        self.y_train = self.dataset.y_train
-        self.y_test = self.dataset.y_test
-        self.train_idx_arr = self.dataset.train_idx_arr
-        self.val_idx_arr = self.dataset.val_idx_arr
-        # Logic to set up the experiment
-        self.items = self.dataset.produce_inputs()
-
-        unique_classes = pd.unique(self.dataframe['event_category'])
-        event_num = len(unique_classes)
-        # droprate = 0.3
-        vector_size = self.dataset.drugs_df.shape[0]
-
-        print("Building the experiment with the following settings:")
-        print(
-            f"Name: {self.experiment_name}, Dataset: {self.dataset}, Model: {self.model}")
-        # Implement additional build logic as needed
-        return self
-
-    def run(self, model_func, batch_size=128, epochs=100):
-        mlflow.set_tracking_uri(self.tracking_uri)
-
-        if mlflow.get_experiment_by_name(self.experiment_name) == None:
-            mlflow.create_experiment(self.experiment_name)
-            mlflow.set_experiment_tags(self.experiment_tags)
-        mlflow.set_experiment(self.experiment_name)
-
-        y_test_label = self.items[0][4]
-        multi_modal = TFMultiModal(
-            model_func=model_func, batch_size=batch_size,  epochs=epochs)  # 100
-        multi_modal.set_data(
-            self.items, self.train_idx_arr, self.val_idx_arr, y_test_label)
-        pred, self.single_results = multi_modal.predict(self.combinations)
-        return self