aiagents4pharma 1.4.3__tar.gz → 1.5.4__tar.gz

{aiagents4pharma-1.4.3 → aiagents4pharma-1.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.4.3
+Version: 1.5.4
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -10,13 +10,16 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: copasi_basico==0.78
 Requires-Dist: coverage==7.6.4
+Requires-Dist: gdown==5.2.0
+Requires-Dist: huggingface_hub==0.26.5
+Requires-Dist: joblib==1.4.2
 Requires-Dist: langchain==0.3.7
 Requires-Dist: langchain-community==0.3.5
 Requires-Dist: langchain-core==0.3.15
 Requires-Dist: langchain-experimental==0.3.3
 Requires-Dist: langchain-openai==0.2.5
 Requires-Dist: matplotlib==3.9.2
-Requires-Dist: openai==1.55.3
+Requires-Dist: openai==1.59.4
 Requires-Dist: pandas==2.2.3
 Requires-Dist: plotly==5.24.1
 Requires-Dist: pydantic==2.9.2
@@ -24,6 +27,7 @@ Requires-Dist: pylint==3.3.1
 Requires-Dist: pytest==8.3.3
 Requires-Dist: streamlit==1.39.0
 Requires-Dist: tabulate==0.9.0
+Requires-Dist: torch==2.5.1
 Requires-Dist: tqdm==4.66.6
 Requires-Dist: mkdocs==1.6.1
 Requires-Dist: mkdocs-jupyter==0.25.1

aiagents4pharma-1.5.4/aiagents4pharma/__init__.py

@@ -0,0 +1,7 @@
+'''
+This file is used to import aiagents4pharma modules.
+'''
+
+from . import talk2biomodels
+from . import talk2cells
+from . import talk2knowledgegraphs

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/__init__.py

@@ -0,0 +1,4 @@
+'''
+This file is used to import the datasets, utils, and tools.
+'''
+from . import datasets

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/datasets/__init__.py

@@ -0,0 +1,7 @@
+'''
+This file is used to import all the models in the package.
+'''
+from . import dataset
+from . import primekg
+from . import starkqa_primekg
+from . import biobridge_primekg

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py

@@ -0,0 +1,562 @@
+"""
+Class for loading BioBridgePrimeKG dataset.
+"""
+
+import os
+import pickle
+import json
+import requests
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from .dataset import Dataset
+from .primekg import PrimeKG
+
+class BioBridgePrimeKG(Dataset):
+    """
+    Class for loading BioBridgePrimeKG dataset.
+    It downloads the data from the BioBridge repo and stores it in the local directory.
+    The data is then loaded into pandas DataFrame of nodes and edges.
+    This class was adapted from the BioBridge repo:
+    https://github.com/RyanWangZf/BioBridge
+    """
+
+    def __init__(self,
+                 primekg_dir: str = "../../../data/primekg/",
+                 local_dir: str = "../../../data/biobridge_primekg/",
+                 random_seed: int=0,
+                 n_neg_samples: int=5):
+        """
+        Constructor for BioBridgePrimeKG class.
+
+        Args:
+            primekg_dir (str): The directory of PrimeKG dataset.
+            local_dir (str): The directory to store the downloaded data.
+            random_seed (int): The random seed value.
+        """
+        self.name: str = "biobridge_primekg"
+        self.primekg_dir: str = primekg_dir
+        self.local_dir: str = local_dir
+        self.random_seed = random_seed
+        self.n_neg_samples = n_neg_samples
+        # Preselected node types:
+        # protein, molecular function, cellular component, biological process, drug, disease
+        self.preselected_node_types = ["protein", "mf", "cc", "bp", "drug", "disease"]
+        self.node_type_map = {
+            "protein": "gene/protein",
+            "mf": "molecular_function",
+            "cc": "cellular_component",
+            "bp": "biological_process",
+            "drug": "drug",
+            "disease": "disease",
+        }
+
+        # Attributes to store the data
+        self.primekg = None
+        self.primekg_triplets = None
+        self.primekg_triplets_negative = None
+        self.data_config = None
+        self.emb_dict = None
+        self.df_train = None
+        self.df_node_train = None
+        self.df_test = None
+        self.df_node_test = None
+        self.node_info_dict = None
+
+        # Set up the dataset
+        self.setup()
+
+    def setup(self):
+        """
+        A method to set up the dataset.
+        """
+        # Make the directories if it doesn't exist
+        os.makedirs(os.path.dirname(self.primekg_dir), exist_ok=True)
+        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)
+
+        # Set the random seed
+        self.set_random_seed(self.random_seed)
+
+        # Set SettingWithCopyWarning  warnings to none
+        pd.options.mode.chained_assignment = None
+
+    def _load_primekg(self) -> PrimeKG:
+        """
+        Private method to load related files of PrimeKG dataset.
+
+        Returns:
+            The PrimeKG dataset.
+        """
+        primekg_data = PrimeKG(local_dir=self.primekg_dir)
+        primekg_data.load_data()
+
+        return primekg_data
+
+    def _download_file(self,
+                       remote_url:str,
+                       local_dir: str,
+                       local_filename: str):
+        """
+        A helper function to download a file from remote URL to the local directory.
+
+        Args:
+            remote_url (str): The remote URL of the file to be downloaded.
+            local_dir (str): The local directory to store the downloaded file.
+            local_filename (str): The local filename to store the downloaded file.
+        """
+        # Make the local directory if it does not exist
+        if not os.path.exists(local_dir):
+            os.makedirs(local_dir)
+        # Download the file from remote URL to local directory
+        local_path = os.path.join(local_dir, local_filename)
+        if os.path.exists(local_path):
+            print(f"File {local_filename} already exists in {local_dir}.")
+        else:
+            print(f"Downloading {local_filename} from {remote_url} to {local_dir}...")
+            response = requests.get(remote_url, stream=True, timeout=300)
+            response.raise_for_status()
+            progress_bar = tqdm(
+                total=int(response.headers.get("content-length", 0)),
+                unit="iB",
+                unit_scale=True,
+            )
+            with open(os.path.join(local_dir, local_filename), "wb") as file:
+                for data in response.iter_content(1024):
+                    progress_bar.update(len(data))
+                    file.write(data)
+            progress_bar.close()
+
+    def _load_data_config(self) -> dict:
+        """
+        Load the data config file of BioBridgePrimeKG dataset.
+
+        Returns:
+            The data config file of BioBridgePrimeKG dataset.
+        """
+        # Download the data config file of BioBridgePrimeKG
+        self._download_file(
+            remote_url= ('https://raw.githubusercontent.com/RyanWangZf/BioBridge/'
+                         'refs/heads/main/data/BindData/data_config.json'),
+            local_dir=self.local_dir,
+            local_filename='data_config.json')
+
+        # Load the downloaded data config file
+        with open(os.path.join(self.local_dir, 'data_config.json'), 'r', encoding='utf-8') as f:
+            data_config = json.load(f)
+
+        return data_config
+
+    def _build_node_embeddings(self) -> dict:
+        """
+        Build the node embeddings for BioBridgePrimeKG dataset.
+
+        Returns:
+            The dictionary of node embeddings.
+        """
+        processed_file_path = os.path.join(self.local_dir, "embeddings", "embedding_dict.pkl")
+        if os.path.exists(processed_file_path):
+            # Load the embeddings from the local directory
+            with open(processed_file_path, "rb") as f:
+                emb_dict_all = pickle.load(f)
+        else:
+            # Download the embeddings from the BioBridge repo and further process them
+            # List of embedding source files
+            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
+                   'refs/heads/main/data/embeddings/esm2b_unimo_pubmedbert/')
+            file_list = [f"{n}.pkl" for n in self.preselected_node_types]
+
+            # Download the embeddings
+            for file in file_list:
+                self._download_file(remote_url=os.path.join(url, file),
+                                    local_dir=os.path.join(self.local_dir, "embeddings"),
+                                    local_filename=file)
+
+            # Unified embeddings
+            emb_dict_all = {}
+            for file in file_list:
+                with open(os.path.join(self.local_dir, "embeddings", file), "rb") as f:
+                    emb = pickle.load(f)
+                emb_ar = emb["embedding"]
+                if not isinstance(emb_ar, list):
+                    emb_ar = emb_ar.tolist()
+                emb_dict_all.update(dict(zip(emb["node_index"], emb_ar)))
+
+            # Store embeddings
+            with open(processed_file_path, "wb") as f:
+                pickle.dump(emb_dict_all, f)
+
+        return emb_dict_all
+
+    def _build_full_triplets(self) -> tuple[pd.DataFrame, dict]:
+        """
+        Build the full triplets for BioBridgePrimeKG dataset.
+
+        Returns:
+            The full triplets for BioBridgePrimeKG dataset.
+            The dictionary of node information.
+        """
+        processed_file_path = os.path.join(self.local_dir, "processed", "triplet_full.tsv.gz")
+        if os.path.exists(processed_file_path):
+            # Load the file from the local directory
+            with open(processed_file_path, "rb") as f:
+                primekg_triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            # Load each dataframe in the local directory
+            node_info_dict = {}
+            for i, node_type in enumerate(self.preselected_node_types):
+                with open(os.path.join(self.local_dir, "processed",
+                                       f"{node_type}.csv"), "rb") as f:
+                    df_node = pd.read_csv(f)
+                node_info_dict[self.node_type_map[node_type]] = df_node
+        else:
+            # Download the related files from the BioBridge repo and further process them
+            # List of processed files
+            url = ('https://media.githubusercontent.com/media/RyanWangZf/BioBridge/'
+                   'refs/heads/main/data/Processed/')
+            file_list = ["protein", "molecular", "cellular", "biological", "drug", "disease"]
+
+            # Download the processed files
+            for i, file in enumerate(file_list):
+                self._download_file(remote_url=os.path.join(url, f"{file}.csv"),
+                                    local_dir=os.path.join(self.local_dir, "processed"),
+                                    local_filename=f"{self.preselected_node_types[i]}.csv")
+
+            # Build the node index list
+            node_info_dict = {}
+            node_index_list = []
+            for i, file in enumerate(file_list):
+                df_node = pd.read_csv(os.path.join(self.local_dir, "processed",
+                                                   f"{self.preselected_node_types[i]}.csv"))
+                node_info_dict[self.node_type_map[self.preselected_node_types[i]]] = df_node
+                node_index_list.extend(df_node["node_index"].tolist())
+
+            # Filter the PrimeKG dataset to take into account only the selected node types
+            primekg_triplets = self.primekg.get_edges().copy()
+            primekg_triplets = primekg_triplets[
+                primekg_triplets["head_index"].isin(node_index_list) &\
+                primekg_triplets["tail_index"].isin(node_index_list)
+            ]
+            primekg_triplets = primekg_triplets.reset_index(drop=True)
+
+            # Perform mapping of node types
+            primekg_triplets["head_type"] = primekg_triplets["head_type"].apply(
+                lambda x: self.data_config["node_type"][x]
+            )
+            primekg_triplets["tail_type"] = primekg_triplets["tail_type"].apply(
+                lambda x: self.data_config["node_type"][x]
+            )
+
+            # Perform mapping of relation types
+            primekg_triplets["display_relation"] = primekg_triplets["display_relation"].apply(
+                lambda x: self.data_config["relation_type"][x]
+            )
+
+            # Store the processed triplets
+            primekg_triplets.to_csv(processed_file_path, sep="\t", compression="gzip", index=False)
+
+        return primekg_triplets, node_info_dict
+
+    def _build_train_test_split(self) -> tuple[pd.DataFrame, pd.DataFrame,
+                                               pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+        """
+        Build the train-test split for BioBridgePrimeKG dataset.
+
+        Returns:
+            The train triplets for BioBridgePrimeKG dataset.
+            The train nodes for BioBridgePrimeKG dataset.
+            The test triplets for BioBridgePrimeKG dataset.
+            The test nodes for BioBridgePrimeKG dataset.
+            The full triplets for BioBridgePrimeKG dataset.
+        """
+        if os.path.exists(os.path.join(self.local_dir, "processed",
+                                       "triplet_full_altered.tsv.gz")):
+            # Load each dataframe in the local directory
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_train.tsv.gz"), "rb") as f:
+                df_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "node_train.tsv.gz"), "rb") as f:
+                df_node_train = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_test.tsv.gz"), "rb") as f:
+                df_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "node_test.tsv.gz"), "rb") as f:
+                df_node_test = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+
+            with open(os.path.join(self.local_dir, "processed",
+                                   "triplet_full_altered.tsv.gz"), "rb") as f:
+                triplets = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+        else:
+            # Filtering out some nodes in the embedding dictionary
+            triplets = self.primekg_triplets.copy()
+            triplets = triplets[
+                triplets["head_index"].isin(list(self.emb_dict.keys())) &\
+                triplets["tail_index"].isin(list(self.emb_dict.keys()))
+            ].reset_index(drop=True)
+
+            # Perform splitting of the triplets
+            list_split = {
+                "train": [],
+                "test": [],
+            }
+            node_split = {
+                "train": {
+                    "node_index": [],
+                    "node_type": [],
+                },
+                "test": {
+                    "node_index": [],
+                    "node_type": [],
+                }
+            }
+            # Loop over the node types
+            for node_type in triplets["head_type"].unique():
+                df_sub = triplets[triplets["head_type"] == node_type]
+                all_x_indexes = df_sub["head_index"].unique()
+                # By default, we use 90% of the nodes for training and 10% for testing
+                te_x_indexes = np.random.choice(
+                    all_x_indexes, size=int(0.1*len(all_x_indexes)), replace=False
+                )
+                df_subs = {}
+                df_subs["test"] = df_sub[df_sub["head_index"].isin(te_x_indexes)]
+                df_subs["train"] = df_sub[~df_sub["head_index"].isin(te_x_indexes)]
+                list_split["train"].append(df_subs["train"])
+                list_split["test"].append(df_subs["test"])
+
+                # record the split
+                node_index = {}
+                node_index["train"] = df_subs["train"]["head_index"].unique()
+                node_split["train"]["node_index"].extend(node_index["train"].tolist())
+                node_split["train"]["node_type"].extend([node_type]*len(node_index["train"]))
+                node_index["test"] = df_subs["test"]["head_index"].unique()
+                node_split["test"]["node_index"].extend(node_index["test"].tolist())
+                node_split["test"]["node_type"].extend([node_type]*len(node_index["test"]))
+
+                print(f"Number of {node_type} nodes in train: {len(node_index['train'])}")
+                print(f"Number of {node_type} nodes in test: {len(node_index['test'])}")
+
+            # Prepare train and test DataFrames
+            df_train = pd.concat(list_split["train"])
+            df_node_train = pd.DataFrame(node_split["train"])
+            df_test = pd.concat(list_split["test"])
+            df_node_test = pd.DataFrame(node_split["test"])
+
+            # Store each dataframe in the local directory
+            df_train.to_csv(os.path.join(self.local_dir, "processed", "triplet_train.tsv.gz"),
+                            sep="\t", compression="gzip", index=False)
+            df_node_train.to_csv(os.path.join(self.local_dir, "processed", "node_train.tsv.gz"),
+                                sep="\t", compression="gzip", index=False)
+            df_test.to_csv(os.path.join(self.local_dir, "processed", "triplet_test.tsv.gz"),
+                           sep="\t", compression="gzip", index=False)
+            df_node_test.to_csv(os.path.join(self.local_dir, "processed", "node_test.tsv.gz"),
+                                sep="\t", compression="gzip", index=False)
+            # Store altered full triplets as well
+            triplets.to_csv(os.path.join(self.local_dir, "processed",
+                                         "triplet_full_altered.tsv.gz"),
+                            sep="\t", compression="gzip", index=False)
+
+        return df_train, df_node_train, df_test, df_node_test, triplets
+
+    # def _negative_sampling(self,
+    #                        batch_df: pd.DataFrame,
+    #                        process_index: int,
+    #                        index_map: dict,
+    #                        node_train_dict: dict) -> pd.DataFrame:
+    #     """
+    #     A helper function to perform negative sampling for a batch of triplets.
+    #     """
+    #     negative_y_index_list = []
+    #     for _, row in tqdm(batch_df.iterrows(),
+    #                        total=batch_df.shape[0],
+    #                        desc=f"Process {process_index}"):
+    #         x_index = row['head_index']
+    #         # y_index = row['y_index']
+    #         y_index_type = row['tail_type']
+    #         paired_y_index_list = index_map[x_index]
+
+    #         # sample a list of negative y_index
+    #         node_train_sub = node_train_dict[y_index_type]
+    #         negative_y_index = node_train_sub[
+    #             ~node_train_sub['node_index'].isin(paired_y_index_list)
+    #         ]['node_index'].sample(self.n_neg_samples).tolist()
+    #         negative_y_index_list.append(negative_y_index)
+
+    #     batch_df.loc[:, 'negative_tail_index'] = negative_y_index_list
+    #     return batch_df
+
+    # def _build_negative_triplets(self,
+    #                              chunk_size: int=100000,
+    #                              n_neg_samples: int=10):
+    #     """
+    #     Build the negative triplets for BioBridgePrimeKG dataset.
+    #     """
+    #     processed_file_path = os.path.join(self.local_dir,
+    #                                        "processed",
+    #                                        "triplet_train_negative.tsv.gz")
+    #     if os.path.exists(processed_file_path):
+    #         # Load the negative triplets from the local directory
+    #         with open(processed_file_path, "rb") as f:
+    #             triplets_negative = pd.read_csv(f, sep="\t", compression="gzip", low_memory=False)
+    #     else:
+    #         # Set the number samples for negative sampling
+    #         self.n_neg_samples = n_neg_samples
+
+    #         # Split node list by type
+    #         node_train_dict = {}
+    #         type_list = self.df_node_train['node_type'].unique()
+    #         for node_type in type_list:
+    #             node_train_dict[node_type] = self.df_node_train[
+    #                 self.df_node_train['node_type'] == node_type
+    #             ].reset_index(drop=True)
+
+    #         # create an index mapping from x_index to y_index
+    #         index_map = self.df_train[
+    #             ['head_index', 'tail_index']
+    #         ].drop_duplicates().groupby('head_index').agg(list).to_dict()['tail_index']
+
+    #         # Negative sampling
+    #         batch_df_list = []
+    #         for i in tqdm(range(0, self.df_train.shape[0], chunk_size)):
+    #             batch_df_list.append(self.df_train.iloc[i:i+chunk_size])
+    #         # Process negative sampling
+    #         results = [
+    #             self._negative_sampling(batch_df,
+    #                                     num_piece,
+    #                                     index_map,
+    #                                     node_train_dict)
+    #                                     for num_piece, batch_df in enumerate(batch_df_list)
+    #         ]
+
+    #         # Store the negative triplets
+    #         triplets_negative = pd.concat(results, axis=0)
+    #         triplets_negative.to_csv(processed_file_path,
+    #                                  sep="\t", compression="gzip", index=False)
+
+    #     # Set attribute
+    #     self.primekg_triplets_negative = triplets_negative
+
+    #     return triplets_negative
+
+    # def load_data(self,
+    #               build_neg_triplest: bool= False,
+    #               chunk_size: int=100000,
+    #               n_neg_samples: int=10):
+
+    def load_data(self):
+        """
+        Load the BioBridgePrimeKG dataset into pandas DataFrame of nodes and edges.
+
+        Args:
+            build_neg_triplest (bool): Whether to build negative triplets.
+            chunk_size (int): The chunk size for negative sampling.
+            n_neg_samples (int): The number of negative samples for negative sampling.
+        """
+        # Load PrimeKG dataset
+        print("Loading PrimeKG dataset...")
+        self.primekg = self._load_primekg()
+
+        # Load data config file of BioBridgePrimeKG
+        print("Loading data config file of BioBridgePrimeKG...")
+        self.data_config = self._load_data_config()
+
+        # Build node embeddings
+        print("Building node embeddings...")
+        self.emb_dict = self._build_node_embeddings()
+
+        # Build full triplets
+        print("Building full triplets...")
+        self.primekg_triplets, self.node_info_dict = self._build_full_triplets()
+
+        # Build train-test split
+        print("Building train-test split...")
+        self.df_train, self.df_node_train, self.df_test, self.df_node_test, self.primekg_triplets =\
+        self._build_train_test_split()
+
+        # if build_neg_triplest:
+        #     # Build negative triplets
+        #     print("Building negative triplets...")
+        #     self.primekg_triplets_negative = self._build_negative_triplets(
+        #         chunk_size=chunk_size,
+        #         n_neg_samples=n_neg_samples
+        #     )
+
+    def set_random_seed(self, seed: int):
+        """
+        Set the random seed for reproducibility.
+
+        Args:
+            seed (int): The random seed value.
+        """
+        np.random.seed(seed)
+
+    def get_primekg(self) -> PrimeKG:
+        """
+        Get the PrimeKG dataset.
+
+        Returns:
+            The PrimeKG dataset.
+        """
+        return self.primekg
+
+    def get_data_config(self) -> dict:
+        """
+        Get the data config file of BioBridgePrimeKG dataset.
+
+        Returns:
+            The data config file of BioBridgePrimeKG dataset.
+        """
+        return self.data_config
+
+    def get_node_embeddings(self) -> dict:
+        """
+        Get the node embeddings for BioBridgePrimeKG dataset.
+
+        Returns:
+            The dictionary of node embeddings.
+        """
+        return self.emb_dict
+
+    def get_primekg_triplets(self) -> pd.DataFrame:
+        """
+        Get the full triplets for BioBridgePrimeKG dataset.
+
+        Returns:
+            The full triplets for BioBridgePrimeKG dataset.
+        """
+        return self.primekg_triplets
+
+    # def get_primekg_triplets_negative(self) -> pd.DataFrame:
+    #     """
+    #     Get the negative triplets for BioBridgePrimeKG dataset.
+
+    #     Returns:
+    #         The negative triplets for BioBridgePrimeKG dataset.
+    #     """
+    #     return self.primekg_triplets_negative
+
+    def get_train_test_split(self) -> dict:
+        """
+        Get the train-test split for BioBridgePrimeKG dataset.
+
+        Returns:
+            The train-test split for BioBridgePrimeKG dataset.
+        """
+        return {
+            "train": self.df_train,
+            "node_train": self.df_node_train,
+            "test": self.df_test,
+            "node_test": self.df_node_test
+        }
+
+    def get_node_info_dict(self) -> dict:
+        """
+        Get the node information dictionary for BioBridgePrimeKG dataset.
+
+        Returns:
+            The node information dictionary for BioBridgePrimeKG dataset.
+        """
+        return self.node_info_dict

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/datasets/dataset.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+
+"""
+Abstract class for dataset.
+"""
+
+from abc import ABC, abstractmethod
+
+class Dataset(ABC):
+    """
+    Abstract class for dataset.
+    """
+    @abstractmethod
+    def setup(self):
+        """
+        A method to set up the dataset.
+        """
+
+    @abstractmethod
+    def load_data(self):
+        """
+        A method to load the dataset and potentially preprocess it.
+        """

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/datasets/primekg.py

@@ -0,0 +1,201 @@
+"""
+Class for loading PrimeKG dataset.
+"""
+
+import os
+import requests
+from tqdm import tqdm
+import pandas as pd
+from .dataset import Dataset
+
+class PrimeKG(Dataset):
+    """
+    Class for loading PrimeKG dataset.
+    It downloads the data from the Harvard Dataverse and stores it in the local directory.
+    The data is then loaded into pandas DataFrame of nodes and edges.
+    """
+
+    def __init__(self, local_dir: str = "../../../data/primekg/"):
+        """
+        Constructor for PrimeKG class.
+
+        Args:
+            local_dir (str): The local directory where the data will be stored.
+        """
+        self.name: str = "primekg"
+        self.server_path: str = "https://dataverse.harvard.edu/api/access/datafile/"
+        self.file_ids: dict = {"nodes": 6180617, "edges": 6180616}
+        self.local_dir: str = local_dir
+
+        # Attributes to store the data
+        self.nodes: pd.DataFrame = None
+        self.edges: pd.DataFrame = None
+
+        # Set up the dataset
+        self.setup()
+
+    def setup(self):
+        """
+        A method to set up the dataset.
+        """
+        # Make the directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)
+
+
+    def _download_file(self, remote_url:str, local_path: str):
+        """
+        A helper function to download a file from remote URL to the local directory.
+
+        Args:
+            remote_url (str): The remote URL of the file to be downloaded.
+            local_path (str): The local path where the file will be saved.
+        """
+        response = requests.get(remote_url, stream=True, timeout=300)
+        response.raise_for_status()
+        progress_bar = tqdm(
+            total=int(response.headers.get("content-length", 0)),
+            unit="iB",
+            unit_scale=True,
+        )
+        with open(local_path, "wb") as file:
+            for data in response.iter_content(1024):
+                progress_bar.update(len(data))
+                file.write(data)
+        progress_bar.close()
+
+    def _load_nodes(self) -> pd.DataFrame:
+        """
+        Private method to load the nodes dataframe of PrimeKG dataset.
+        This method downloads the nodes file from the Harvard Dataverse if it does not exist
+        in the local directory. Otherwise, it loads the data from the local directory.
+        It further processes the dataframe of nodes and returns it.
+
+        Returns:
+            The nodes dataframe of PrimeKG dataset.
+        """
+        local_file = os.path.join(self.local_dir, f"{self.name}_nodes.tsv.gz")
+        if os.path.exists(local_file):
+            print(f"{local_file} already exists. Loading the data from the local directory.")
+
+            # Load the dataframe from the local directory and assign it to the nodes attribute
+            nodes = pd.read_csv(local_file, sep="\t", compression="gzip", low_memory=False)
+        else:
+            print(f"Downloading node file from {self.server_path}{self.file_ids['nodes']}")
+
+            # Download the file from the Harvard Dataverse with designated file_id for node
+            self._download_file(f"{self.server_path}{self.file_ids['nodes']}",
+                                os.path.join(self.local_dir, "nodes.tab"))
+
+            # Load the downloaded file into a pandas DataFrame
+            nodes = pd.read_csv(os.path.join(self.local_dir, "nodes.tab"),
+                                     sep="\t", low_memory=False)
+
+            # Further processing of the dataframe
+            nodes = nodes[
+                ["node_index", "node_name", "node_source", "node_id", "node_type"]
+            ]
+
+            # Store compressed dataframe in the local directory
+            nodes.to_csv(local_file, index=False, sep="\t", compression="gzip")
+
+        return nodes
+
+    def _load_edges(self, nodes: pd.DataFrame) -> pd.DataFrame:
+        """
+        Private method to load the edges dataframe of PrimeKG dataset.
+        This method downloads the edges file from the Harvard Dataverse if it does not exist
+        in the local directory. Otherwise, it loads the data from the local directory.
+        It further processes the dataframe of edges and returns it.
+
+        Args:
+            nodes (pd.DataFrame): The nodes dataframe of PrimeKG dataset.
+
+        Returns:
+            The edges dataframe of PrimeKG dataset.
+        """
+        local_file = os.path.join(self.local_dir, f"{self.name}_edges.tsv.gz")
+        if os.path.exists(local_file):
+            print(f"{local_file} already exists. Loading the data from the local directory.")
+
+            # Load the dataframe from the local directory and assign it to the edges attribute
+            edges = pd.read_csv(local_file, sep="\t", compression="gzip", low_memory=False)
+        else:
+            print(f"Downloading edge file from {self.server_path}{self.file_ids['edges']}")
+
+            # Download the file from the Harvard Dataverse with designated file_id for edge
+            self._download_file(f"{self.server_path}{self.file_ids['edges']}",
+                                os.path.join(self.local_dir, "edges.csv"))
+
+            # Load the downloaded file into a pandas DataFrame
+            edges = pd.read_csv(os.path.join(self.local_dir, "edges.csv"),
+                                     sep=",", low_memory=False)
+
+            # Further processing of the dataframe
+            edges = edges.merge(
+                nodes, left_on="x_index", right_on="node_index"
+            )
+            edges.drop(["x_index"], axis=1, inplace=True)
+            edges.rename(
+                columns={
+                    "node_index": "head_index",
+                    "node_name": "head_name",
+                    "node_source": "head_source",
+                    "node_id": "head_id",
+                    "node_type": "head_type",
+                },
+                inplace=True,
+            )
+            edges = edges.merge(
+                nodes, left_on="y_index", right_on="node_index"
+            )
+            edges.drop(["y_index"], axis=1, inplace=True)
+            edges.rename(
+                columns={
+                    "node_index": "tail_index",
+                    "node_name": "tail_name",
+                    "node_source": "tail_source",
+                    "node_id": "tail_id",
+                    "node_type": "tail_type"
+                },
+                inplace=True,
+            )
+            edges = edges[
+                [
+                    "head_index", "head_name", "head_source", "head_id", "head_type",
+                    "tail_index", "tail_name", "tail_source", "tail_id", "tail_type",
+                    "display_relation", "relation",
+                ]
+            ]
+
+            # Store compressed dataframe in the local directory
+            edges.to_csv(local_file, index=False, sep="\t", compression="gzip")
+
+        return edges
+
+    def load_data(self):
+        """
+        Load the PrimeKG dataset into pandas DataFrame of nodes and edges.
+        """
+        print("Loading nodes of PrimeKG dataset ...")
+        self.nodes = self._load_nodes()
+
+        print("Loading edges of PrimeKG dataset ...")
+        self.edges = self._load_edges(self.nodes)
+
+    def get_nodes(self) -> pd.DataFrame:
+        """
+        Get the nodes dataframe of PrimeKG dataset.
+
+        Returns:
+            The nodes dataframe of PrimeKG dataset.
+        """
+        return self.nodes
+
+    def get_edges(self) -> pd.DataFrame:
+        """
+        Get the edges dataframe of PrimeKG dataset.
+
+        Returns:
+            The edges dataframe of PrimeKG dataset.
+        """
+        return self.edges

aiagents4pharma-1.5.4/aiagents4pharma/talk2knowledgegraphs/datasets/starkqa_primekg.py

@@ -0,0 +1,201 @@
+"""
+Class for loading StarkQAPrimeKG dataset.
+"""
+
+import os
+import shutil
+import pickle
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import torch
+from huggingface_hub import hf_hub_download, list_repo_files
+import gdown
+from .dataset import Dataset
+
+class StarkQAPrimeKG(Dataset):
+    """
+    Class for loading StarkQAPrimeKG dataset.
+    It downloads the data from the HuggingFace repo and stores it in the local directory.
+    The data is then loaded into pandas DataFrame of QA pairs, dictionary of split indices,
+    and node information.
+    """
+
+    def __init__(self, local_dir: str = "../../../data/starkqa_primekg/"):
+        """
+        Constructor for StarkQAPrimeKG class.
+
+        Args:
+            local_dir (str): The local directory to store the dataset files.
+        """
+        self.name: str = "starkqa_primekg"
+        self.hf_repo_id: str = "snap-stanford/stark"
+        self.local_dir: str = local_dir
+        # Attributes to store the data
+        self.starkqa: pd.DataFrame = None
+        self.starkqa_split_idx: dict = None
+        self.starkqa_node_info: dict = None
+        self.query_emb_dict: dict = None
+        self.node_emb_dict: dict = None
+
+        # Set up the dataset
+        self.setup()
+
+    def setup(self):
+        """
+        A method to set up the dataset.
+        """
+        # Make the directory if it doesn't exist
+        os.makedirs(os.path.dirname(self.local_dir), exist_ok=True)
+
+    def _load_stark_repo(self) -> tuple[pd.DataFrame, dict, dict]:
+        """
+        Private method to load related files of StarkQAPrimeKG dataset.
+
+        Returns:
+            The nodes dataframe of StarkQAPrimeKG dataset.
+            The split indices of StarkQAPrimeKG dataset.
+            The node information of StarkQAPrimeKG dataset.
+        """
+        # Download the file if it does not exist in the local directory
+        # Otherwise, load the data from the local directory
+        local_file = os.path.join(self.local_dir, "qa/prime/stark_qa/stark_qa.csv")
+        if os.path.exists(local_file):
+            print(f"{local_file} already exists. Loading the data from the local directory.")
+        else:
+            print(f"Downloading files from {self.hf_repo_id}")
+
+            # List all related files in the HuggingFace Hub repository
+            files = list_repo_files(self.hf_repo_id, repo_type="dataset")
+            files = [f for f in files if ((f.startswith("qa/prime/") or
+                                           f.startswith("skb/prime/")) and f.find("raw") == -1)]
+
+            # Download and save each file in the specified folder
+            for file in tqdm(files):
+                _ = hf_hub_download(self.hf_repo_id,
+                                    file,
+                                    repo_type="dataset",
+                                    local_dir=self.local_dir)
+
+            # Unzip the processed files
+            shutil.unpack_archive(
+                os.path.join(self.local_dir, "skb/prime/processed.zip"),
+                os.path.join(self.local_dir, "skb/prime/")
+            )
+
+        # Load StarkQA dataframe
+        starkqa = pd.read_csv(
+            os.path.join(self.local_dir, "qa/prime/stark_qa/stark_qa.csv"),
+            low_memory=False)
+
+        # Read split indices
+        qa_indices = sorted(starkqa['id'].tolist())
+        starkqa_split_idx = {}
+        for split in ['train', 'val', 'test', 'test-0.1']:
+            indices_file = os.path.join(self.local_dir, "qa/prime/split", f'{split}.index')
+            with open(indices_file, 'r', encoding='utf-8') as f:
+                indices = f.read().strip().split('\n')
+            query_ids = [int(idx) for idx in indices]
+            starkqa_split_idx[split] = np.array(
+                [qa_indices.index(query_id) for query_id in query_ids]
+            )
+
+        # Load the node info of PrimeKG preprocessed for StarkQA
+        with open(os.path.join(self.local_dir, 'skb/prime/processed/node_info.pkl'), 'rb') as f:
+            starkqa_node_info = pickle.load(f)
+
+        return starkqa, starkqa_split_idx, starkqa_node_info
+
+    def _load_stark_embeddings(self) -> tuple[dict, dict]:
+        """
+        Private method to load the embeddings of StarkQAPrimeKG dataset.
+
+        Returns:
+            The query embeddings of StarkQAPrimeKG dataset.
+            The node embeddings of StarkQAPrimeKG dataset.
+        """
+        # Load the provided embeddings of query and nodes
+        # Note that they utilized 'text-embedding-ada-002' for embeddings
+        emb_model = 'text-embedding-ada-002'
+        query_emb_url = 'https://drive.google.com/uc?id=1MshwJttPZsHEM2cKA5T13SIrsLeBEdyU'
+        node_emb_url = 'https://drive.google.com/uc?id=16EJvCMbgkVrQ0BuIBvLBp-BYPaye-Edy'
+
+        # Prepare respective directories to store the embeddings
+        emb_dir = os.path.join(self.local_dir, emb_model)
+        query_emb_dir = os.path.join(emb_dir, "query")
+        node_emb_dir = os.path.join(emb_dir, "doc")
+        os.makedirs(query_emb_dir, exist_ok=True)
+        os.makedirs(node_emb_dir, exist_ok=True)
+        query_emb_path = os.path.join(query_emb_dir, "query_emb_dict.pt")
+        node_emb_path = os.path.join(node_emb_dir, "candidate_emb_dict.pt")
+
+        # Download the embeddings if they do not exist in the local directory
+        if not os.path.exists(query_emb_path) or not os.path.exists(node_emb_path):
+            # Download the query embeddings
+            gdown.download(query_emb_url, query_emb_path, quiet=False)
+
+            # Download the node embeddings
+            gdown.download(node_emb_url, node_emb_path, quiet=False)
+
+        # Load the embeddings
+        query_emb_dict = torch.load(query_emb_path)
+        node_emb_dict = torch.load(node_emb_path)
+
+        return query_emb_dict, node_emb_dict
+
+    def load_data(self):
+        """
+        Load the StarkQAPrimeKG dataset into pandas DataFrame of QA pairs,
+        dictionary of split indices, and node information.
+        """
+        print("Loading StarkQAPrimeKG dataset...")
+        self.starkqa, self.starkqa_split_idx, self.starkqa_node_info = self._load_stark_repo()
+
+        print("Loading StarkQAPrimeKG embeddings...")
+        self.query_emb_dict, self.node_emb_dict = self._load_stark_embeddings()
+
+
+    def get_starkqa(self) -> pd.DataFrame:
+        """
+        Get the dataframe of StarkQAPrimeKG dataset, containing the QA pairs.
+
+        Returns:
+            The nodes dataframe of PrimeKG dataset.
+        """
+        return self.starkqa
+
+    def get_starkqa_split_indicies(self) -> dict:
+        """
+        Get the split indices of StarkQAPrimeKG dataset.
+
+        Returns:
+            The split indices of StarkQAPrimeKG dataset.
+        """
+        return self.starkqa_split_idx
+
+    def get_starkqa_node_info(self) -> dict:
+        """
+        Get the node information of StarkQAPrimeKG dataset.
+
+        Returns:
+            The node information of StarkQAPrimeKG dataset.
+        """
+        return self.starkqa_node_info
+
+    def get_query_embeddings(self) -> dict:
+        """
+        Get the query embeddings of StarkQAPrimeKG dataset.
+
+        Returns:
+            The query embeddings of StarkQAPrimeKG dataset.
+        """
+        return self.query_emb_dict
+
+    def get_node_embeddings(self) -> dict:
+        """
+        Get the node embeddings of StarkQAPrimeKG dataset.
+
+        Returns:
+            The node embeddings of StarkQAPrimeKG dataset.
+        """
+        return self.node_emb_dict

{aiagents4pharma-1.4.3 → aiagents4pharma-1.5.4}/aiagents4pharma.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.4.3
+Version: 1.5.4
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -10,13 +10,16 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: copasi_basico==0.78
 Requires-Dist: coverage==7.6.4
+Requires-Dist: gdown==5.2.0
+Requires-Dist: huggingface_hub==0.26.5
+Requires-Dist: joblib==1.4.2
 Requires-Dist: langchain==0.3.7
 Requires-Dist: langchain-community==0.3.5
 Requires-Dist: langchain-core==0.3.15
 Requires-Dist: langchain-experimental==0.3.3
 Requires-Dist: langchain-openai==0.2.5
 Requires-Dist: matplotlib==3.9.2
-Requires-Dist: openai==1.55.3
+Requires-Dist: openai==1.59.4
 Requires-Dist: pandas==2.2.3
 Requires-Dist: plotly==5.24.1
 Requires-Dist: pydantic==2.9.2
@@ -24,6 +27,7 @@ Requires-Dist: pylint==3.3.1
 Requires-Dist: pytest==8.3.3
 Requires-Dist: streamlit==1.39.0
 Requires-Dist: tabulate==0.9.0
+Requires-Dist: torch==2.5.1
 Requires-Dist: tqdm==4.66.6
 Requires-Dist: mkdocs==1.6.1
 Requires-Dist: mkdocs-jupyter==0.25.1

{aiagents4pharma-1.4.3 → aiagents4pharma-1.5.4}/aiagents4pharma.egg-info/SOURCES.txt

@@ -18,4 +18,10 @@ aiagents4pharma/talk2biomodels/tools/custom_plotter.py
 aiagents4pharma/talk2biomodels/tools/fetch_parameters.py
 aiagents4pharma/talk2biomodels/tools/model_description.py
 aiagents4pharma/talk2biomodels/tools/search_models.py
-aiagents4pharma/talk2biomodels/tools/simulate_model.py
+aiagents4pharma/talk2biomodels/tools/simulate_model.py
+aiagents4pharma/talk2knowledgegraphs/__init__.py
+aiagents4pharma/talk2knowledgegraphs/datasets/__init__.py
+aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py
+aiagents4pharma/talk2knowledgegraphs/datasets/dataset.py
+aiagents4pharma/talk2knowledgegraphs/datasets/primekg.py
+aiagents4pharma/talk2knowledgegraphs/datasets/starkqa_primekg.py

{aiagents4pharma-1.4.3 → aiagents4pharma-1.5.4}/aiagents4pharma.egg-info/requires.txt

@@ -1,12 +1,15 @@
 copasi_basico==0.78
 coverage==7.6.4
+gdown==5.2.0
+huggingface_hub==0.26.5
+joblib==1.4.2
 langchain==0.3.7
 langchain-community==0.3.5
 langchain-core==0.3.15
 langchain-experimental==0.3.3
 langchain-openai==0.2.5
 matplotlib==3.9.2
-openai==1.55.3
+openai==1.59.4
 pandas==2.2.3
 plotly==5.24.1
 pydantic==2.9.2
@@ -14,6 +17,7 @@ pylint==3.3.1
 pytest==8.3.3
 streamlit==1.39.0
 tabulate==0.9.0
+torch==2.5.1
 tqdm==4.66.6
 mkdocs==1.6.1
 mkdocs-jupyter==0.25.1

{aiagents4pharma-1.4.3 → aiagents4pharma-1.5.4}/pyproject.toml

@@ -15,13 +15,16 @@ classifiers = [
 dependencies = [
     "copasi_basico==0.78",
     "coverage==7.6.4",
+    "gdown==5.2.0",
+    "huggingface_hub==0.26.5",
+    "joblib==1.4.2",
     "langchain==0.3.7",
     "langchain-community==0.3.5",
     "langchain-core==0.3.15",
     "langchain-experimental==0.3.3",
     "langchain-openai==0.2.5",
     "matplotlib==3.9.2",
-    "openai==1.55.3",
+    "openai==1.59.4",
     "pandas==2.2.3",
     "plotly==5.24.1",
     "pydantic==2.9.2",
@@ -29,6 +32,7 @@ dependencies = [
     "pytest==8.3.3",
     "streamlit==1.39.0",
     "tabulate==0.9.0",
+    "torch==2.5.1",
     "tqdm==4.66.6",
     "mkdocs==1.6.1",
     "mkdocs-jupyter==0.25.1",
@@ -48,7 +52,9 @@ version = {file = "release_version.txt"}
 packages = ["aiagents4pharma",
             "aiagents4pharma.talk2biomodels",
             "aiagents4pharma.talk2biomodels.models",
-            "aiagents4pharma.talk2biomodels.tools"]
+            "aiagents4pharma.talk2biomodels.tools",
+            "aiagents4pharma.talk2knowledgegraphs",
+            "aiagents4pharma.talk2knowledgegraphs.datasets"]
 
 # [tool.setuptools.packages.find]
 # where = ["aiagents4pharma", "aiagents4pharma.talk2biomodels"]

aiagents4pharma-1.5.4/release_version.txt

@@ -0,0 +1 @@
+v1.5.4

aiagents4pharma-1.4.3/aiagents4pharma/__init__.py

@@ -1,5 +0,0 @@
-'''
-This file is used to import the talk2biomodels module.
-'''
-
-from . import talk2biomodels

aiagents4pharma-1.4.3/release_version.txt

@@ -1 +0,0 @@
-v1.4.3