Update backend_utils.py

Author: puririshi98

torch_geometric/utils/rag/backend_utils.py

@@ -25,7 +25,6 @@
     TripletLike,
 )
 from torch_geometric.data.large_graph_indexer import EDGE_RELATION
-from torch_geometric.datasets.web_qsp_dataset import retrieval_via_pcst
 from torch_geometric.distributed import (
     LocalFeatureStore,
     LocalGraphStore,
@@ -50,6 +49,129 @@ def preprocess_triplet(triplet: TripletLike) -> TripletLike:
     h, r, t = triplet
     return str(h).lower(), str(r).lower(), str(t).lower()
 
+@no_type_check
+def retrieval_via_pcst(
+    data: Data,
+    q_emb: Tensor,
+    textual_nodes: Any,
+    textual_edges: Any,
+    topk: int = 3,
+    topk_e: int = 5,
+    cost_e: float = 0.5,
+    num_clusters: int = 1,
+) -> Tuple[Data, str]:
+
+    # skip PCST for bad graphs
+    booly = data.edge_attr is None or data.edge_attr.numel() == 0
+    booly = booly or data.x is None or data.x.numel() == 0
+    booly = booly or data.edge_index is None or data.edge_index.numel() == 0
+    if not booly:
+        c = 0.01
+
+        from pcst_fast import pcst_fast
+
+        root = -1
+        pruning = 'gw'
+        verbosity_level = 0
+        if topk > 0:
+            n_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.x)
+            topk = min(topk, data.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()
+        else:
+            n_prizes = torch.zeros(data.num_nodes)
+
+        if topk_e > 0:
+            e_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.edge_attr)
+            topk_e = min(topk_e, e_prizes.unique().size(0))
+
+            topk_e_values, _ = torch.topk(e_prizes.unique(), 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 = e_prizes == topk_e_values[k]
+                value = min((topk_e - k) / sum(indices), last_topk_e_value - c)
+                e_prizes[indices] = value
+                last_topk_e_value = value * (1 - c)
+            # reduce the cost of the edges so that at least one edge is chosen
+            cost_e = min(cost_e, e_prizes.max().item() * (1 - c / 2))
+        else:
+            e_prizes = torch.zeros(data.num_edges)
+
+        costs = []
+        edges = []
+        virtual_n_prizes = []
+        virtual_edges = []
+        virtual_costs = []
+        mapping_n = {}
+        mapping_e = {}
+        for i, (src, dst) in enumerate(data.edge_index.t().numpy()):
+            prize_e = e_prizes[i]
+            if prize_e <= cost_e:
+                mapping_e[len(edges)] = i
+                edges.append((src, dst))
+                costs.append(cost_e - prize_e)
+            else:
+                virtual_node_id = data.num_nodes + len(virtual_n_prizes)
+                mapping_n[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 - cost_e)
+
+        prizes = np.concatenate([n_prizes, np.array(virtual_n_prizes)])
+        num_edges = len(edges)
+        if len(virtual_costs) > 0:
+            costs = np.array(costs + virtual_costs)
+            edges = np.array(edges + virtual_edges)
+
+        vertices, edges = pcst_fast(edges, prizes, costs, root, num_clusters,
+                                    pruning, verbosity_level)
+
+        selected_nodes = vertices[vertices < data.num_nodes]
+        selected_edges = [mapping_e[e] for e in edges if e < num_edges]
+        virtual_vertices = vertices[vertices >= data.num_nodes]
+        if len(virtual_vertices) > 0:
+            virtual_vertices = vertices[vertices >= data.num_nodes]
+            virtual_edges = [mapping_n[i] for i in virtual_vertices]
+            selected_edges = np.array(selected_edges + virtual_edges)
+
+        edge_index = data.edge_index[:, selected_edges]
+        selected_nodes = np.unique(
+            np.concatenate(
+                [selected_nodes, edge_index[0].numpy(),
+                 edge_index[1].numpy()]))
+
+        n = textual_nodes.iloc[selected_nodes]
+        e = textual_edges.iloc[selected_edges]
+    else:
+        n = textual_nodes
+        e = textual_edges
+    desc = n.to_csv(index=False) + '\n' + e.to_csv(
+        index=False, columns=['src', 'edge_attr', 'dst'])
+
+    mapping = {n: i for i, n in enumerate(selected_nodes.tolist())}
+    src = [mapping[i] for i in edge_index[0].tolist()]
+    dst = [mapping[i] for i in edge_index[1].tolist()]
+
+    # HACK Added so that the subset of nodes and edges selected can be tracked
+    node_idx = np.array(data.node_idx)[selected_nodes]
+    edge_idx = np.array(data.edge_idx)[selected_edges]
+
+    data = Data(
+        x=data.x[selected_nodes],
+        edge_index=torch.tensor([src, dst]).to(torch.long),
+        edge_attr=data.edge_attr[selected_edges],
+        # HACK Added so that the subset of nodes and edges selected can be tracked
+        node_idx=node_idx,
+        edge_idx=edge_idx,
+    )
+
+    return data, desc
 
 def batch_knn(query_enc: Tensor, embeds: Tensor,
               k: int) -> Iterator[InputNodes]: