aiagents4pharma 1.36.0__py3-none-any.whl → 1.38.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/tools/multimodal_subgraph_extraction.py

@@ -0,0 +1,374 @@
+"""
+Tool for performing multimodal subgraph extraction.
+"""
+
+from typing import Type, Annotated
+import logging
+import pickle
+import numpy as np
+import pandas as pd
+import hydra
+import networkx as nx
+from pydantic import BaseModel, Field
+from langchain_core.tools import BaseTool
+from langchain_core.messages import ToolMessage
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+import torch
+from torch_geometric.data import Data
+from ..utils.extractions.multimodal_pcst import MultimodalPCSTPruning
+from ..utils.embeddings.ollama import EmbeddingWithOllama
+from .load_arguments import ArgumentData
+
+# Initialize logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class MultimodalSubgraphExtractionInput(BaseModel):
+    """
+    MultimodalSubgraphExtractionInput is a Pydantic model representing an input
+    for extracting a subgraph.
+
+    Args:
+        prompt: Prompt to interact with the backend.
+        tool_call_id: Tool call ID.
+        state: Injected state.
+        arg_data: Argument for analytical process over graph data.
+    """
+
+    tool_call_id: Annotated[str, InjectedToolCallId] = Field(
+        description="Tool call ID."
+    )
+    state: Annotated[dict, InjectedState] = Field(description="Injected state.")
+    prompt: str = Field(description="Prompt to interact with the backend.")
+    arg_data: ArgumentData = Field(
+        description="Experiment over graph data.", default=None
+    )
+
+
+class MultimodalSubgraphExtractionTool(BaseTool):
+    """
+    This tool performs subgraph extraction based on user's prompt by taking into account
+    the top-k nodes and edges.
+    """
+
+    name: str = "subgraph_extraction"
+    description: str = "A tool for subgraph extraction based on user's prompt."
+    args_schema: Type[BaseModel] = MultimodalSubgraphExtractionInput
+
+    def _prepare_query_modalities(self,
+                                  prompt_emb: list,
+                                  state: Annotated[dict, InjectedState],
+                                  pyg_graph: Data) -> pd.DataFrame:
+        """
+        Prepare the modality-specific query for subgraph extraction.
+
+        Args:
+            prompt_emb: The embedding of the user prompt in a list.
+            state: The injected state for the tool.
+            pyg_graph: The PyTorch Geometric graph Data.
+
+        Returns:
+            A DataFrame containing the query embeddings and modalities.
+        """
+        # Initialize dataframes
+        multimodal_df = pd.DataFrame({"name": []})
+        query_df = pd.DataFrame({"node_id": [],
+                                 "node_type": [],
+                                 "x": [],
+                                 "desc_x": [],
+                                 "use_description": []})
+
+        # Loop over the uploaded files and find multimodal files
+        for i in range(len(state["uploaded_files"])):
+            # Check if multimodal file is uploaded
+            if state["uploaded_files"][i]["file_type"] == "multimodal":
+                # Read the Excel file
+                multimodal_df = pd.read_excel(state["uploaded_files"][i]["file_path"],
+                                              sheet_name=None)
+
+        # Check if the multimodal_df is empty
+        if len(multimodal_df) > 0:
+            # Merge all obtained dataframes into a single dataframe
+            multimodal_df = pd.concat(multimodal_df).reset_index()
+            multimodal_df.drop(columns=["level_1"], inplace=True)
+            multimodal_df.rename(columns={"level_0": "q_node_type",
+                                        "name": "q_node_name"}, inplace=True)
+            # Since an excel sheet name could not contain a `/`,
+            # but the node type can be 'gene/protein' as exists in the PrimeKG
+            multimodal_df["q_node_type"] = multimodal_df.q_node_type.apply(
+                lambda x: x.replace('-', '/')
+            )
+
+            # Convert PyG graph to a DataFrame for easier filtering
+            graph_df = pd.DataFrame({
+                "node_id": pyg_graph.node_id,
+                "node_name": pyg_graph.node_name,
+                "node_type": pyg_graph.node_type,
+                "x": pyg_graph.x,
+                "desc_x": pyg_graph.desc_x.tolist(),
+            })
+
+            # Make a query dataframe by merging the graph_df and multimodal_df
+            query_df = graph_df.merge(multimodal_df, how='cross')
+            query_df = query_df[
+                query_df.apply(
+                    lambda x:
+                    (x['q_node_name'].lower() in x['node_name'].lower()) & # node name
+                    (x['node_type'] == x['q_node_type']), # node type
+                    axis=1
+                )
+            ]
+            query_df = query_df[['node_id', 'node_type', 'x', 'desc_x']].reset_index(drop=True)
+            query_df['use_description'] = False # set to False for modal-specific embeddings
+
+            # Update the state by adding the the selected node IDs
+            state["selections"] = query_df.groupby("node_type")["node_id"].apply(list).to_dict()
+
+        # Append a user prompt to the query dataframe
+        query_df = pd.concat([
+            query_df,
+            pd.DataFrame({
+                'node_id': 'user_prompt',
+                'node_type': 'prompt',
+                'x': prompt_emb,
+                'desc_x': prompt_emb,
+                'use_description': True # set to True for user prompt embedding
+            })
+        ]).reset_index(drop=True)
+
+        return query_df
+
+    def _perform_subgraph_extraction(self,
+                                     state: Annotated[dict, InjectedState],
+                                     cfg: dict,
+                                     pyg_graph: Data,
+                                     query_df: pd.DataFrame) -> dict:
+        """
+        Perform multimodal subgraph extraction based on modal-specific embeddings.
+
+        Args:
+            state: The injected state for the tool.
+            cfg: The configuration dictionary.
+            pyg_graph: The PyTorch Geometric graph Data.
+            query_df: The DataFrame containing the query embeddings and modalities.
+
+        Returns:
+            A dictionary containing the extracted subgraph with nodes and edges.
+        """
+        # Initialize the subgraph dictionary
+        subgraphs = {}
+        subgraphs["nodes"] = []
+        subgraphs["edges"] = []
+
+        # Loop over query embeddings and modalities
+        for q in query_df.iterrows():
+            # Prepare the PCSTPruning object and extract the subgraph
+            # Parameters were set in the configuration file obtained from Hydra
+            subgraph = MultimodalPCSTPruning(
+                topk=state["topk_nodes"],
+                topk_e=state["topk_edges"],
+                cost_e=cfg.cost_e,
+                c_const=cfg.c_const,
+                root=cfg.root,
+                num_clusters=cfg.num_clusters,
+                pruning=cfg.pruning,
+                verbosity_level=cfg.verbosity_level,
+                use_description=q[1]['use_description'],
+            ).extract_subgraph(pyg_graph,
+                               torch.tensor(q[1]['desc_x']), # description embedding
+                               torch.tensor(q[1]['x']), # modal-specific embedding
+                               q[1]['node_type'])
+
+            # Append the extracted subgraph to the dictionary
+            subgraphs["nodes"].append(subgraph["nodes"].tolist())
+            subgraphs["edges"].append(subgraph["edges"].tolist())
+
+        # Concatenate and get unique node and edge indices
+        subgraphs["nodes"] = np.unique(
+            np.concatenate([np.array(list_) for list_ in subgraphs["nodes"]])
+        )
+        subgraphs["edges"] = np.unique(
+            np.concatenate([np.array(list_) for list_ in subgraphs["edges"]])
+        )
+
+        return subgraphs
+
+    def _prepare_final_subgraph(self,
+                               state:Annotated[dict, InjectedState],
+                               subgraph: dict,
+                               graph: dict,
+                               cfg) -> dict:
+        """
+        Prepare the subgraph based on the extracted subgraph.
+
+        Args:
+            state: The injected state for the tool.
+            subgraph: The extracted subgraph.
+            graph: The initial graph containing PyG and textualized graph.
+            cfg: The configuration dictionary.
+
+        Returns:
+            A dictionary containing the PyG graph, NetworkX graph, and textualized graph.
+        """
+        # print(subgraph)
+        # Prepare the PyTorch Geometric graph
+        mapping = {n: i for i, n in enumerate(subgraph["nodes"].tolist())}
+        pyg_graph = Data(
+            # Node features
+            # x=pyg_graph.x[subgraph["nodes"]],
+            x=[graph["pyg"].x[i] for i in subgraph["nodes"]],
+            node_id=np.array(graph["pyg"].node_id)[subgraph["nodes"]].tolist(),
+            node_name=np.array(graph["pyg"].node_id)[subgraph["nodes"]].tolist(),
+            enriched_node=np.array(graph["pyg"].enriched_node)[subgraph["nodes"]].tolist(),
+            num_nodes=len(subgraph["nodes"]),
+            # Edge features
+            edge_index=torch.LongTensor(
+                [
+                    [
+                        mapping[i]
+                        for i in graph["pyg"].edge_index[:, subgraph["edges"]][0].tolist()
+                    ],
+                    [
+                        mapping[i]
+                        for i in graph["pyg"].edge_index[:, subgraph["edges"]][1].tolist()
+                    ],
+                ]
+            ),
+            edge_attr=graph["pyg"].edge_attr[subgraph["edges"]],
+            edge_type=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            relation=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            label=np.array(graph["pyg"].edge_type)[subgraph["edges"]].tolist(),
+            enriched_edge=np.array(graph["pyg"].enriched_edge)[subgraph["edges"]].tolist(),
+        )
+
+        # Networkx DiGraph construction to be visualized in the frontend
+        nx_graph = nx.DiGraph()
+        # Add nodes with attributes
+        node_colors = {n: cfg.node_colors_dict[k]
+                       for k, v in state["selections"].items() for n in v}
+        for n in pyg_graph.node_name:
+            nx_graph.add_node(n, color=node_colors.get(n, None))
+
+        # Add edges with attributes
+        edges = zip(
+            pyg_graph.edge_index[0].tolist(),
+            pyg_graph.edge_index[1].tolist(),
+            pyg_graph.edge_type
+        )
+        for src, dst, edge_type in edges:
+            nx_graph.add_edge(
+                pyg_graph.node_name[src],
+                pyg_graph.node_name[dst],
+                relation=edge_type,
+                label=edge_type,
+            )
+
+        # Prepare the textualized subgraph
+        textualized_graph = (
+            graph["text"]["nodes"].iloc[subgraph["nodes"]].to_csv(index=False)
+            + "\n"
+            + graph["text"]["edges"].iloc[subgraph["edges"]].to_csv(index=False)
+        )
+
+        return {
+            "graph_pyg": pyg_graph,
+            "graph_nx": nx_graph,
+            "graph_text": textualized_graph,
+        }
+
+    def _run(
+        self,
+        tool_call_id: Annotated[str, InjectedToolCallId],
+        state: Annotated[dict, InjectedState],
+        prompt: str,
+        arg_data: ArgumentData = None,
+    ) -> Command:
+        """
+        Run the subgraph extraction tool.
+
+        Args:
+            tool_call_id: The tool call ID for the tool.
+            state: Injected state for the tool.
+            prompt: The prompt to interact with the backend.
+            arg_data (ArgumentData): The argument data.
+
+        Returns:
+            Command: The command to be executed.
+        """
+        logger.log(logging.INFO, "Invoking subgraph_extraction tool")
+
+        # Load hydra configuration
+        with hydra.initialize(version_base=None, config_path="../configs"):
+            cfg = hydra.compose(
+                config_name="config", overrides=["tools/multimodal_subgraph_extraction=default"]
+            )
+            cfg = cfg.tools.multimodal_subgraph_extraction
+
+        # Retrieve source graph from the state
+        initial_graph = {}
+        initial_graph["source"] = state["dic_source_graph"][-1]  # The last source graph as of now
+        # logger.log(logging.INFO, "Source graph: %s", source_graph)
+
+        # Load the knowledge graph
+        with open(initial_graph["source"]["kg_pyg_path"], "rb") as f:
+            initial_graph["pyg"] = pickle.load(f)
+        with open(initial_graph["source"]["kg_text_path"], "rb") as f:
+            initial_graph["text"] = pickle.load(f)
+
+        # Prepare the query embeddings and modalities
+        query_df = self._prepare_query_modalities(
+            [EmbeddingWithOllama(model_name=cfg.ollama_embeddings[0]).embed_query(prompt)],
+            state,
+            initial_graph["pyg"]
+        )
+
+        # Perform subgraph extraction
+        subgraphs = self._perform_subgraph_extraction(state,
+                                                      cfg,
+                                                      initial_graph["pyg"],
+                                                      query_df)
+
+        # Prepare subgraph as a NetworkX graph and textualized graph
+        final_subgraph = self._prepare_final_subgraph(state,
+                                                      subgraphs,
+                                                      initial_graph,
+                                                      cfg)
+
+        # Prepare the dictionary of extracted graph
+        dic_extracted_graph = {
+            "name": arg_data.extraction_name,
+            "tool_call_id": tool_call_id,
+            "graph_source": initial_graph["source"]["name"],
+            "topk_nodes": state["topk_nodes"],
+            "topk_edges": state["topk_edges"],
+            "graph_dict": {
+                "nodes": list(final_subgraph["graph_nx"].nodes(data=True)),
+                "edges": list(final_subgraph["graph_nx"].edges(data=True)),
+            },
+            "graph_text": final_subgraph["graph_text"],
+            "graph_summary": None,
+        }
+
+        # Prepare the dictionary of updated state
+        dic_updated_state_for_model = {}
+        for key, value in {
+            "dic_extracted_graph": [dic_extracted_graph],
+        }.items():
+            if value:
+                dic_updated_state_for_model[key] = value
+
+        # Return the updated state of the tool
+        return Command(
+            update=dic_updated_state_for_model | {
+                # update the message history
+                "messages": [
+                    ToolMessage(
+                        content=f"Subgraph Extraction Result of {arg_data.extraction_name}",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )

aiagents4pharma/talk2knowledgegraphs/utils/extractions/__init__.py

@@ -2,3 +2,4 @@
 This file is used to import all the models in the package.
 '''
 from . import pcst
+from . import multimodal_pcst

aiagents4pharma/talk2knowledgegraphs/utils/extractions/multimodal_pcst.py

@@ -0,0 +1,292 @@
+"""
+Exctraction of multimodal subgraph using Prize-Collecting Steiner Tree (PCST) algorithm.
+"""
+
+from typing import Tuple, NamedTuple
+import numpy as np
+import pandas as pd
+import torch
+import pcst_fast
+from torch_geometric.data.data import Data
+
+class MultimodalPCSTPruning(NamedTuple):
+    """
+    Prize-Collecting Steiner Tree (PCST) pruning algorithm implementation inspired by G-Retriever
+    (He et al., 'G-Retriever: Retrieval-Augmented Generation for Textual Graph Understanding and
+    Question Answering', NeurIPS 2024) paper.
+    https://arxiv.org/abs/2402.07630
+    https://github.com/XiaoxinHe/G-Retriever/blob/main/src/dataset/utils/retrieval.py
+
+    Args:
+        topk: The number of top nodes to consider.
+        topk_e: The number of top edges to consider.
+        cost_e: The cost of the edges.
+        c_const: The constant value for the cost of the edges computation.
+        root: The root node of the subgraph, -1 for unrooted.
+        num_clusters: The number of clusters.
+        pruning: The pruning strategy to use.
+        verbosity_level: The verbosity level.
+    """
+    topk: int = 3
+    topk_e: int = 3
+    cost_e: float = 0.5
+    c_const: float = 0.01
+    root: int = -1
+    num_clusters: int = 1
+    pruning: str = "gw"
+    verbosity_level: int = 0
+    use_description: bool = False
+
+    def _compute_node_prizes(self,
+                             graph: Data,
+                             query_emb: torch.Tensor,
+                             modality: str) :
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction based on the node type.
+
+        Returns:
+            The prizes of the nodes.
+        """
+        # Convert PyG graph to a DataFrame
+        graph_df = pd.DataFrame({
+            "node_type": graph.node_type,
+            "desc_x": [x.tolist() for x in graph.desc_x],
+            "x": [list(x) for x in graph.x],
+            "score": [0.0 for _ in range(len(graph.node_id))],
+        })
+
+        # Calculate cosine similarity for text features and update the score
+        if self.use_description:
+            graph_df.loc[:, "score"] = torch.nn.CosineSimilarity(dim=-1)(
+                    query_emb,
+                    torch.tensor(list(graph_df.desc_x.values)) # Using textual description features
+                ).tolist()
+        else:
+            graph_df.loc[graph_df["node_type"] == modality,
+                         "score"] = torch.nn.CosineSimilarity(dim=-1)(
+                    query_emb,
+                    torch.tensor(list(graph_df[graph_df["node_type"]== modality].x.values))
+                ).tolist()
+
+        # Set the prizes for nodes based on the similarity scores
+        n_prizes = torch.tensor(graph_df.score.values, dtype=torch.float32)
+        # n_prizes = torch.nn.CosineSimilarity(dim=-1)(query_emb, graph.x)
+        topk = min(self.topk, graph.num_nodes)
+        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
+        n_prizes = torch.zeros_like(n_prizes)
+        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
+
+        return n_prizes
+
+    def _compute_edge_prizes(self,
+                             graph: Data,
+                             text_emb: torch.Tensor) :
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+
+        Returns:
+            The prizes of the nodes.
+        """
+        # Note that as of now, the edge features are based on textual features
+        # Compute prizes for edges
+        e_prizes = torch.nn.CosineSimilarity(dim=-1)(text_emb, graph.edge_attr)
+        unique_prizes, inverse_indices = e_prizes.unique(return_inverse=True)
+        topk_e = min(self.topk_e, unique_prizes.size(0))
+        topk_e_values, _ = torch.topk(unique_prizes, topk_e, largest=True)
+        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
+        last_topk_e_value = topk_e
+        for k in range(topk_e):
+            indices = inverse_indices == (
+                unique_prizes == topk_e_values[k]
+            ).nonzero(as_tuple=True)[0]
+            value = min((topk_e - k) / indices.sum().item(), last_topk_e_value)
+            e_prizes[indices] = value
+            last_topk_e_value = value * (1 - self.c_const)
+
+        return e_prizes
+
+    def compute_prizes(self,
+                       graph: Data,
+                       text_emb: torch.Tensor,
+                       query_emb: torch.Tensor,
+                       modality: str):
+        """
+        Compute the node prizes based on the cosine similarity between the query and nodes,
+        as well as the edge prizes based on the cosine similarity between the query and edges.
+        Note that the node and edge embeddings shall use the same embedding model and dimensions
+        with the query.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction based on node type.
+
+        Returns:
+            The prizes of the nodes and edges.
+        """
+        # Compute prizes for nodes
+        n_prizes = self._compute_node_prizes(graph, query_emb, modality)
+
+        # Compute prizes for edges
+        e_prizes = self._compute_edge_prizes(graph, text_emb)
+
+        return {"nodes": n_prizes, "edges": e_prizes}
+
+    def compute_subgraph_costs(self,
+                               graph: Data,
+                               prizes: dict) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Compute the costs in constructing the subgraph proposed by G-Retriever paper.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            prizes: The prizes of the nodes and the edges.
+
+        Returns:
+            edges: The edges of the subgraph, consisting of edges and number of edges without
+                virtual edges.
+            prizes: The prizes of the subgraph.
+            costs: The costs of the subgraph.
+        """
+        # Logic to reduce the cost of the edges such that at least one edge is selected
+        updated_cost_e = min(
+            self.cost_e,
+            prizes["edges"].max().item() * (1 - self.c_const / 2),
+        )
+
+        # Initialize variables
+        edges = []
+        costs = []
+        virtual = {
+            "n_prizes": [],
+            "edges": [],
+            "costs": [],
+        }
+        mapping = {"nodes": {}, "edges": {}}
+
+        # Compute the costs, edges, and virtual variables based on the prizes
+        for i, (src, dst) in enumerate(graph.edge_index.T.numpy()):
+            prize_e = prizes["edges"][i]
+            if prize_e <= updated_cost_e:
+                mapping["edges"][len(edges)] = i
+                edges.append((src, dst))
+                costs.append(updated_cost_e - prize_e)
+            else:
+                virtual_node_id = graph.num_nodes + len(virtual["n_prizes"])
+                mapping["nodes"][virtual_node_id] = i
+                virtual["edges"].append((src, virtual_node_id))
+                virtual["edges"].append((virtual_node_id, dst))
+                virtual["costs"].append(0)
+                virtual["costs"].append(0)
+                virtual["n_prizes"].append(prize_e - updated_cost_e)
+        prizes = np.concatenate([prizes["nodes"], np.array(virtual["n_prizes"])])
+        edges_dict = {}
+        edges_dict["edges"] = edges
+        edges_dict["num_prior_edges"] = len(edges)
+        # Final computation of the costs and edges based on the virtual costs and virtual edges
+        if len(virtual["costs"]) > 0:
+            costs = np.array(costs + virtual["costs"])
+            edges = np.array(edges + virtual["edges"])
+            edges_dict["edges"] = edges
+
+        return edges_dict, prizes, costs, mapping
+
+    def get_subgraph_nodes_edges(
+        self, graph: Data, vertices: np.ndarray, edges_dict: dict, mapping: dict,
+    ) -> dict:
+        """
+        Get the selected nodes and edges of the subgraph based on the vertices and edges computed
+        by the PCST algorithm.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            vertices: The vertices of the subgraph computed by the PCST algorithm.
+            edges_dict: The dictionary of edges of the subgraph computed by the PCST algorithm,
+                and the number of prior edges (without virtual edges).
+            mapping: The mapping dictionary of the nodes and edges.
+            num_prior_edges: The number of edges before adding virtual edges.
+
+        Returns:
+            The selected nodes and edges of the extracted subgraph.
+        """
+        # Get edges information
+        edges = edges_dict["edges"]
+        num_prior_edges = edges_dict["num_prior_edges"]
+        # Retrieve the selected nodes and edges based on the given vertices and edges
+        subgraph_nodes = vertices[vertices < graph.num_nodes]
+        subgraph_edges = [mapping["edges"][e] for e in edges if e < num_prior_edges]
+        virtual_vertices = vertices[vertices >= graph.num_nodes]
+        if len(virtual_vertices) > 0:
+            virtual_vertices = vertices[vertices >= graph.num_nodes]
+            virtual_edges = [mapping["nodes"][i] for i in virtual_vertices]
+            subgraph_edges = np.array(subgraph_edges + virtual_edges)
+        edge_index = graph.edge_index[:, subgraph_edges]
+        subgraph_nodes = np.unique(
+            np.concatenate(
+                [subgraph_nodes, edge_index[0].numpy(), edge_index[1].numpy()]
+            )
+        )
+
+        return {"nodes": subgraph_nodes, "edges": subgraph_edges}
+
+    def extract_subgraph(self,
+                         graph: Data,
+                         text_emb: torch.Tensor,
+                         query_emb: torch.Tensor,
+                         modality: str) -> dict:
+        """
+        Perform the Prize-Collecting Steiner Tree (PCST) algorithm to extract the subgraph.
+
+        Args:
+            graph: The knowledge graph in PyTorch Geometric Data format.
+            text_emb: The textual description embedding in PyTorch Tensor format.
+            query_emb: The query embedding in PyTorch Tensor format. This can be an embedding of
+                a prompt, sequence, or any other feature to be used for the subgraph extraction.
+            modality: The modality to use for the subgraph extraction
+                (e.g., "text", "sequence", "smiles").
+
+        Returns:
+            The selected nodes and edges of the subgraph.
+        """
+        # Assert the topk and topk_e values for subgraph retrieval
+        assert self.topk > 0, "topk must be greater than or equal to 0"
+        assert self.topk_e > 0, "topk_e must be greater than or equal to 0"
+
+        # Retrieve the top-k nodes and edges based on the query embedding
+        prizes = self.compute_prizes(graph, text_emb, query_emb, modality)
+
+        # Compute costs in constructing the subgraph
+        edges_dict, prizes, costs, mapping = self.compute_subgraph_costs(
+            graph, prizes
+        )
+
+        # Retrieve the subgraph using the PCST algorithm
+        result_vertices, result_edges = pcst_fast.pcst_fast(
+            edges_dict["edges"],
+            prizes,
+            costs,
+            self.root,
+            self.num_clusters,
+            self.pruning,
+            self.verbosity_level,
+        )
+
+        subgraph = self.get_subgraph_nodes_edges(
+            graph,
+            result_vertices,
+            {"edges": result_edges, "num_prior_edges": edges_dict["num_prior_edges"]},
+            mapping)
+
+        return subgraph

aiagents4pharma/talk2scholars/configs/tools/zotero_read/default.yaml

@@ -2,6 +2,7 @@
 library_type: "user" # Type of library ('user' or 'group')
 default_limit: 2
 request_timeout: 10
+chunk_size: 16384  # Size (in bytes) for streaming PDF download chunks
 user_id: ${oc.env:ZOTERO_USER_ID} # Load from environment variable
 api_key: ${oc.env:ZOTERO_API_KEY} # Load from environment variable