Spaces:
Sleeping
Sleeping
| """ | |
| Convert a persisted LlamaIndex PropertyGraphIndex into a PyG Data checkpoint. | |
| Run from the repository root: | |
| python -m src.graph.pyg_converter | |
| """ | |
| import torch | |
| from torch_geometric.data import Data | |
| from sentence_transformers import SentenceTransformer | |
| from llama_index.core import StorageContext, load_index_from_storage | |
| from src.utils import UTF8LocalFileSystem | |
| def convert_to_pyg( | |
| persist_dir: str = "./storage_graph", | |
| output_path: str = "./storage_graph/pyg_data.pt", | |
| ): | |
| print(f"Loading graph from {persist_dir}…") | |
| try: | |
| storage_context = StorageContext.from_defaults( | |
| persist_dir=persist_dir, fs=UTF8LocalFileSystem() | |
| ) | |
| index = load_index_from_storage(storage_context) | |
| pg_store = index.property_graph_store | |
| except Exception as e: | |
| print(f"Error loading index: {e}") | |
| return | |
| nodes_dict = pg_store.graph.nodes | |
| rel_dict = ( | |
| pg_store.graph.relations.values() | |
| if isinstance(pg_store.graph.relations, dict) | |
| else pg_store.graph.relations | |
| ) | |
| node_id_to_idx: dict = {} | |
| node_texts: list[str] = [] | |
| for i, (node_id, node_data) in enumerate(nodes_dict.items()): | |
| node_id_to_idx[node_id] = i | |
| text_rep = str(node_id) | |
| if hasattr(node_data, "properties") and node_data.properties: | |
| text_rep = str( | |
| node_data.properties.get("name") | |
| or node_data.properties.get("id") | |
| or node_id | |
| ) | |
| if hasattr(node_data, "text") and node_data.text: | |
| text_rep += " " + node_data.text | |
| node_texts.append(text_rep) | |
| print(f"Extracted {len(node_texts)} nodes.") | |
| print("Encoding nodes with BAAI/bge-small-en-v1.5…") | |
| embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5") | |
| embeddings = embed_model.encode(node_texts, show_progress_bar=True, convert_to_tensor=True) | |
| src_indices: list[int] = [] | |
| dst_indices: list[int] = [] | |
| for rel in rel_dict: | |
| if hasattr(rel, "source_id") and hasattr(rel, "target_id"): | |
| src, dst = rel.source_id, rel.target_id | |
| else: | |
| src = rel.get("source_id") | |
| dst = rel.get("target_id") | |
| if src in node_id_to_idx and dst in node_id_to_idx: | |
| src_indices.append(node_id_to_idx[src]) | |
| dst_indices.append(node_id_to_idx[dst]) | |
| print(f"Extracted {len(src_indices)} edges.") | |
| edge_index = torch.tensor([src_indices, dst_indices], dtype=torch.long) | |
| data = Data(x=embeddings.cpu(), edge_index=edge_index) | |
| idx_to_node_id = {v: k for k, v in node_id_to_idx.items()} | |
| print(f"Saving PyG checkpoint to {output_path}…") | |
| torch.save( | |
| { | |
| "pyg_data": data, | |
| "node_id_to_idx": node_id_to_idx, | |
| "idx_to_node_id": idx_to_node_id, | |
| }, | |
| output_path, | |
| ) | |
| print("Done.") | |
| if __name__ == "__main__": | |
| convert_to_pyg() | |