aiagents4pharma 1.41.0__py3-none-any.whl → 1.43.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/milvus_data_dump.py

@@ -1,438 +1,796 @@
-# pylint: disable=wrong-import-position
 #!/usr/bin/env python3
+# pylint: skip-file
 """
-Script to load PrimeKG multimodal data into Milvus database.
-This script runs after Milvus container is ready and loads the .pkl file data.
+Dynamic Cross-Platform PrimeKG Multimodal Data Loader for Milvus Database.
+Automatically detects system capabilities and chooses appropriate libraries and configurations.
+
+Supports:
+- Windows, Linux, macOS
+- CPU-only mode (pandas/numpy)
+- NVIDIA GPU mode (cudf/cupy)
+- Dynamic index selection based on hardware
+- Automatic dependency installation
 """
 
-import os
-import sys
-import subprocess
 import glob
 import logging
-from typing import Dict, Any, List
-
-def install_packages():
-    """Install required packages."""
-    packages = [
-        "pip install --extra-index-url=https://pypi.nvidia.com cudf-cu12",
-        "pip install --extra-index-url=https://pypi.nvidia.com dask-cudf-cu12",
-        "pip install pymilvus==2.5.11",
-        "pip install numpy==1.26.4",
-        "pip install pandas==2.1.3",
-        "pip install tqdm==4.67.1",
-    ]
-
-    print("[DATA LOADER] Installing required packages...")
-    for package_cmd in packages:
-        print(f"[DATA LOADER] Running: {package_cmd}")
-        result = subprocess.run(package_cmd.split(), capture_output=True, text=True, check=True)
-        if result.returncode != 0:
-            print(f"[DATA LOADER] Error installing package: {result.stderr}")
-            sys.exit(1)
-    print("[DATA LOADER] All packages installed successfully!")
-
-# Install packages first
-install_packages()
-
-try:
-    import cudf
-    import cupy as cp
-except ImportError as e:
-    print("[DATA LOADER] cudf or cupy not found. Please ensure they are installed correctly.")
-    sys.exit(1)
-
-from pymilvus import (
-    db,
-    connections,
-    FieldSchema,
-    CollectionSchema,
-    DataType,
-    Collection,
-    utility
-)
-from tqdm import tqdm
+import os
+import platform
+import subprocess
+import sys
+from typing import Any, Dict, List, Optional, Union
 
 # Configure logging
-logging.basicConfig(level=logging.INFO, format='[DATA LOADER] %(message)s')
+logging.basicConfig(level=logging.INFO, format="[DATA LOADER] %(message)s")
 logger = logging.getLogger(__name__)
 
-class MilvusDataLoader:
-    """
-    Class to handle loading of BioBridge-PrimeKG multimodal data into Milvus.
-    """
+
+class SystemDetector:
+    """Detect system capabilities and choose appropriate libraries."""
+
+    def __init__(self):
+        self.os_type = platform.system().lower()  # 'windows', 'linux', 'darwin'
+        self.architecture = platform.machine().lower()  # 'x86_64', 'arm64', etc.
+        self.has_nvidia_gpu = self._detect_nvidia_gpu()
+        self.use_gpu = (
+            self.has_nvidia_gpu and self.os_type != "darwin"
+        )  # No CUDA on macOS
+
+        logger.info("System Detection Results:")
+        logger.info("  OS: %s", self.os_type)
+        logger.info("  Architecture: %s", self.architecture)
+        logger.info("  NVIDIA GPU detected: %s", self.has_nvidia_gpu)
+        logger.info("  Will use GPU acceleration: %s", self.use_gpu)
+
+    def _detect_nvidia_gpu(self) -> bool:
+        """Detect if NVIDIA GPU is available."""
+        try:
+            # Try nvidia-smi command
+            result = subprocess.run(
+                ["nvidia-smi"], capture_output=True, text=True, timeout=10
+            )
+            return result.returncode == 0
+        except (
+            subprocess.TimeoutExpired,
+            FileNotFoundError,
+            subprocess.SubprocessError,
+        ):
+            return False
+
+    def get_required_packages(self) -> List[str]:
+        """Get list of packages to install based on system capabilities - matches original logic."""
+        if self.use_gpu and self.os_type == "linux":
+            # Exact package list from original script for GPU mode
+            packages = [
+                # "pip install --extra-index-url=https://pypi.nvidia.com cudf-cu12",
+                # "pip install --extra-index-url=https://pypi.nvidia.com dask-cudf-cu12",
+                "pip install pymilvus==2.5.11",
+                "pip install numpy==1.26.4",
+                "pip install pandas==2.1.3",
+                "pip install tqdm==4.67.1",
+            ]
+            return packages
+        else:
+            # CPU-only packages
+            packages = [
+                "pip install pymilvus==2.5.11",
+                "pip install numpy==1.26.4",
+                "pip install pandas==2.1.3",
+                "pip install tqdm==4.67.1",
+            ]
+            return packages
+
+    def install_packages(self):
+        """Install required packages using original script's exact logic."""
+        packages = self.get_required_packages()
+
+        logger.info(
+            "Installing packages for %s system%s",
+            self.os_type,
+            " with GPU support" if self.use_gpu else "",
+        )
+
+        for package_cmd in packages:
+            logger.info("Running: %s", package_cmd)
+            try:
+                result = subprocess.run(
+                    package_cmd.split(),
+                    capture_output=True,
+                    text=True,
+                    check=True,
+                    timeout=300,
+                )
+                if result.returncode != 0:
+                    logger.error("Error installing package: %s", result.stderr)
+                    if "cudf" in package_cmd or "dask-cudf" in package_cmd:
+                        logger.warning(
+                            "GPU package installation failed, falling back to CPU mode"
+                        )
+                        self.use_gpu = False
+                        return self.install_packages()  # Retry with CPU packages
+                    else:
+                        sys.exit(1)
+                else:
+                    logger.info("Successfully installed: %s", package_cmd.split()[-1])
+            except subprocess.CalledProcessError as e:
+                logger.error("Failed to install %s: %s", package_cmd, e.stderr)
+                if "cudf" in package_cmd:
+                    logger.warning(
+                        "GPU package installation failed, falling back to CPU mode"
+                    )
+                    self.use_gpu = False
+                    return self.install_packages()  # Retry with CPU packages
+                else:
+                    raise
+            except subprocess.TimeoutExpired:
+                logger.error("Installation timeout for package: %s", package_cmd)
+                raise
+
+
+class DynamicDataLoader:
+    """Dynamic data loader that adapts to system capabilities."""
+
     def __init__(self, config: Dict[str, Any]):
-        """Initialize the MilvusDataLoader with configuration parameters."""
+        """Initialize with system detection and dynamic library loading."""
         self.config = config
-        self.milvus_host = config.get('milvus_host', 'localhost')
-        self.milvus_port = config.get('milvus_port', '19530')
-        self.milvus_user = config.get('milvus_user', 'root')
-        self.milvus_password = config.get('milvus_password', 'Milvus')
-        self.milvus_database = config.get('milvus_database', 't2kg_primekg')
-        self.data_dir = config.get('data_dir',
-                                   'tests/files/biobridge_multimodal/')
-        self.batch_size = config.get('batch_size', 500)
-        self.chunk_size = config.get('chunk_size', 5)
-
-    def normalize_matrix(self, m, axis=1):
-        """Normalize each row of a 2D matrix using CuPy."""
-        norms = cp.linalg.norm(m, axis=axis, keepdims=True)
-        return m / norms
-
-    def normalize_vector(self, v):
-        """Normalize a vector using CuPy."""
-        v = cp.asarray(v)
-        norm = cp.linalg.norm(v)
-        return v / norm
+        self.detector = SystemDetector()
+
+        # Install packages if needed
+        if config.get("auto_install_packages", True):
+            self.detector.install_packages()
+
+        # Import libraries based on system capabilities
+        self._import_libraries()
+
+        # Configuration - exact original parameters
+        self.milvus_host = config.get("milvus_host", "localhost")
+        self.milvus_port = config.get("milvus_port", "19530")
+        self.milvus_user = config.get("milvus_user", "root")
+        self.milvus_password = config.get("milvus_password", "Milvus")
+        self.milvus_database = config.get("milvus_database", "t2kg_primekg")
+        self.data_dir = config.get("data_dir", "./data")
+        self.batch_size = config.get("batch_size", 500)
+        self.chunk_size = config.get("chunk_size", 5)  # Original chunk_size parameter
+
+        # Dynamic settings based on hardware
+        self.use_gpu = self.detector.use_gpu
+        self.normalize_vectors = self.use_gpu  # Only normalize for GPU (original logic)
+        self.vector_index_type = "GPU_CAGRA" if self.use_gpu else "HNSW"
+        self.metric_type = "IP" if self.use_gpu else "COSINE"
+
+        logger.info("Loader Configuration:")
+        logger.info("  Using GPU acceleration: %s", self.use_gpu)
+        logger.info("  Vector normalization: %s", self.normalize_vectors)
+        logger.info("  Vector index type: %s", self.vector_index_type)
+        logger.info("  Metric type: %s", self.metric_type)
+        logger.info("  Data directory: %s", self.data_dir)
+        logger.info("  Batch size: %s", self.batch_size)
+        logger.info("  Chunk size: %s", self.chunk_size)
+
+    def _import_libraries(self):
+        """Dynamically import libraries - matches original script's import logic."""
+        # Always import base libraries
+        import numpy as np
+        import pandas as pd
+        from pymilvus import (
+            Collection,
+            CollectionSchema,
+            DataType,
+            FieldSchema,
+            connections,
+            db,
+            utility,
+        )
+        from tqdm import tqdm
+
+        self.pd = pd
+        self.np = np
+        self.tqdm = tqdm
+        self.pymilvus_modules = {
+            "db": db,
+            "connections": connections,
+            "FieldSchema": FieldSchema,
+            "CollectionSchema": CollectionSchema,
+            "DataType": DataType,
+            "Collection": Collection,
+            "utility": utility,
+        }
+
+        # Conditionally import GPU libraries - matches original error handling
+        if self.detector.use_gpu:
+            try:
+                import cudf  # pyright: ignore
+                import cupy as cp  # pyright: ignore
+
+                self.cudf = cudf
+                self.cp = cp
+                logger.info("Successfully imported GPU libraries (cudf, cupy)")
+            except ImportError as e:
+                logger.error(
+                    "[DATA LOADER] cudf or cupy not found. Please ensure they are installed correctly."
+                )
+                logger.error("Import error: %s", str(e))
+                # Match original script's exit behavior for critical GPU import failure
+                if not os.getenv("FORCE_CPU", "false").lower() == "true":
+                    logger.error(
+                        "GPU libraries required but not available. Set FORCE_CPU=true to use CPU mode."
+                    )
+                    sys.exit(1)
+                else:
+                    logger.warning("Falling back to CPU mode due to FORCE_CPU=true")
+                    self.detector.use_gpu = False
+                    self.use_gpu = False
+
+    def _read_dataframe(
+        self, file_path: str, columns: Optional[List[str]] = None
+    ) -> Union["pd.DataFrame", "cudf.DataFrame"]:  # type: ignore[reportUndefinedVariable]  # noqa: F821
+        """Read dataframe using appropriate library."""
+        if self.use_gpu:
+            return self.cudf.read_parquet(file_path, columns=columns)
+        else:
+            return self.pd.read_parquet(file_path, columns=columns)
+
+    def _concat_dataframes(
+        self, df_list: List, ignore_index: bool = True
+    ) -> Union["pd.DataFrame", "cudf.DataFrame"]:  # type: ignore[reportUndefinedVariable]  # noqa: F821
+        """Concatenate dataframes using appropriate library."""
+        if self.use_gpu:
+            return self.cudf.concat(df_list, ignore_index=ignore_index)
+        else:
+            return self.pd.concat(df_list, ignore_index=ignore_index)
+
+    def _normalize_matrix(self, matrix, axis: int = 1):
+        """Normalize matrix using appropriate library."""
+        if not self.normalize_vectors:
+            return matrix
+
+        if self.use_gpu:
+            # Use cupy for GPU
+            matrix_cp = self.cp.asarray(matrix).astype(self.cp.float32)
+            norms = self.cp.linalg.norm(matrix_cp, axis=axis, keepdims=True)
+            return matrix_cp / norms
+        else:
+            # Use numpy for CPU (but we don't normalize for CPU/COSINE)
+            return matrix
+
+    def _extract_embeddings(self, df, column_name: str):
+        """Extract embeddings and convert to appropriate format."""
+        if self.use_gpu:
+            # cuDF list extraction
+            emb_data = self.cp.asarray(df[column_name].list.leaves).astype(
+                self.cp.float32
+            )
+            return emb_data.reshape(df.shape[0], -1)
+        else:
+            # pandas extraction
+            emb_list = []
+            for emb in df[column_name]:
+                if isinstance(emb, list):
+                    emb_list.append(emb)
+                else:
+                    emb_list.append(emb.tolist() if hasattr(emb, "tolist") else emb)
+            return self.np.array(emb_list, dtype=self.np.float32)
+
+    def _to_list(self, data):
+        """Convert data to list format for Milvus insertion."""
+        if self.use_gpu:
+            # For cuDF data, use to_arrow().to_pylist()
+            if hasattr(data, "to_arrow"):
+                return data.to_arrow().to_pylist()
+            elif hasattr(data, "tolist"):
+                # Fallback for cupy arrays
+                return data.tolist()
+            else:
+                return list(data)
+        else:
+            # For pandas/numpy data
+            if hasattr(data, "tolist"):
+                return data.tolist()
+            elif hasattr(data, "to_arrow"):
+                return data.to_arrow().to_pylist()
+            else:
+                return list(data)
 
     def connect_to_milvus(self):
         """Connect to Milvus and setup database."""
         logger.info("Connecting to Milvus at %s:%s", self.milvus_host, self.milvus_port)
 
-        connections.connect(
+        self.pymilvus_modules["connections"].connect(
             alias="default",
             host=self.milvus_host,
             port=self.milvus_port,
             user=self.milvus_user,
-            password=self.milvus_password
+            password=self.milvus_password,
         )
 
         # Check if database exists, create if it doesn't
-        if self.milvus_database not in db.list_database():
+        if self.milvus_database not in self.pymilvus_modules["db"].list_database():
             logger.info("Creating database: %s", self.milvus_database)
-            db.create_database(self.milvus_database)
+            self.pymilvus_modules["db"].create_database(self.milvus_database)
 
         # Switch to the desired database
-        db.using_database(self.milvus_database)
+        self.pymilvus_modules["db"].using_database(self.milvus_database)
         logger.info("Using database: %s", self.milvus_database)
 
     def load_graph_data(self):
-        """Load the pickle file containing graph data."""
+        """Load the parquet files containing graph data."""
         logger.info("Loading graph data from: %s", self.data_dir)
 
         if not os.path.exists(self.data_dir):
             raise FileNotFoundError(f"Data directory not found: {self.data_dir}")
 
-        # Load dataframes containing nodes and edges
-        # Loop over nodes and edges
         graph = {}
         for element in ["nodes", "edges"]:
-            # Make an empty dictionary for each folder
             graph[element] = {}
             for stage in ["enrichment", "embedding"]:
-                print(element, stage)
-                # Create the file pattern for the current subfolder
-                file_list = glob.glob(os.path.join(self.data_dir,
-                                                   element,
-                                                   stage,
-                                                   '*.parquet.gzip'))
-                print(file_list)
-                # Read and concatenate all dataframes in the folder
-                # Except the edges embedding, which is too large to read in one go
-                # We are using a chunk size to read the edges embedding in smaller parts instead
+                logger.info("Processing %s %s", element, stage)
+
+                file_list = glob.glob(
+                    os.path.join(self.data_dir, element, stage, "*.parquet.gzip")
+                )
+                logger.info("Found %d files for %s %s", len(file_list), element, stage)
+
+                if not file_list:
+                    logger.warning("No files found for %s %s", element, stage)
+                    continue
+
+                # For edges embedding, process in chunks due to size
                 if element == "edges" and stage == "embedding":
-                    # For edges embedding, only read two columns:
-                    # triplet_index and edge_emb
-                    # Loop by chunks
                     chunk_size = self.chunk_size
                     graph[element][stage] = []
                     for i in range(0, len(file_list), chunk_size):
-                        chunk_files = file_list[i:i+chunk_size]
-                        chunk_df = cudf.concat([
-                            cudf.read_parquet(f, columns=["triplet_index", "edge_emb"])
-                            for f in chunk_files
-                        ], ignore_index=True)
+                        chunk_files = file_list[i : i + chunk_size]
+                        chunk_df_list = []
+                        for f in chunk_files:
+                            df = self._read_dataframe(
+                                f, columns=["triplet_index", "edge_emb"]
+                            )
+                            chunk_df_list.append(df)
+                        chunk_df = self._concat_dataframes(
+                            chunk_df_list, ignore_index=True
+                        )
                         graph[element][stage].append(chunk_df)
                 else:
-                    # For nodes and edges enrichment,
-                    # read and concatenate all dataframes in the folder
-                    # This includes the nodes embedding,
-                    # which is small enough to read in one go
-                    graph[element][stage] = cudf.concat([
-                        cudf.read_parquet(f) for f in file_list
-                    ], ignore_index=True)
+                    # For other combinations, read all files
+                    df_list = []
+                    for f in file_list:
+                        df = self._read_dataframe(f)
+                        df_list.append(df)
+                    graph[element][stage] = self._concat_dataframes(
+                        df_list, ignore_index=True
+                    )
 
         logger.info("Graph data loaded successfully")
         return graph
 
-    def create_nodes_collection(self, nodes_df: cudf.DataFrame):
+    def _get_embedding_dimension(self, df, column_name: str) -> int:
+        """Get embedding dimension using original script's exact logic."""
+        first_emb = df.iloc[0][column_name]
+        if self.use_gpu:
+            # cuDF format - matches original: len(nodes_df.iloc[0]['desc_emb'].to_arrow().to_pylist()[0])
+            return len(first_emb.to_arrow().to_pylist()[0])
+        else:
+            # pandas format
+            if isinstance(first_emb, list):
+                return len(first_emb)
+            else:
+                return len(
+                    first_emb.tolist() if hasattr(first_emb, "tolist") else first_emb
+                )
+
+    def create_nodes_collection(self, nodes_df):
         """Create and populate the main nodes collection."""
         logger.info("Creating main nodes collection...")
         node_coll_name = f"{self.milvus_database}_nodes"
 
+        # Get embedding dimension
+        emb_dim = self._get_embedding_dimension(nodes_df, "desc_emb")
+
         node_fields = [
-            FieldSchema(name="node_index",
-                        dtype=DataType.INT64,
-                        is_primary=True),
-            FieldSchema(name="node_id",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024),
-            FieldSchema(name="node_name",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024,
-                        enable_analyzer=True,
-                        enable_match=True),
-            FieldSchema(name="node_type",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024,
-                        enable_analyzer=True,
-                        enable_match=True),
-            FieldSchema(name="desc",
-                        dtype=DataType.VARCHAR,
-                        max_length=40960,
-                        enable_analyzer=True,
-                        enable_match=True),
-            FieldSchema(name="desc_emb",
-                        dtype=DataType.FLOAT_VECTOR,
-                        dim=len(nodes_df.iloc[0]['desc_emb'].to_arrow().to_pylist()[0])),
+            self.pymilvus_modules["FieldSchema"](
+                name="node_index",
+                dtype=self.pymilvus_modules["DataType"].INT64,
+                is_primary=True,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="node_id",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="node_name",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+                enable_analyzer=True,
+                enable_match=True,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="node_type",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+                enable_analyzer=True,
+                enable_match=True,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="desc",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=40960,
+                enable_analyzer=True,
+                enable_match=True,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="desc_emb",
+                dtype=self.pymilvus_modules["DataType"].FLOAT_VECTOR,
+                dim=emb_dim,
+            ),
         ]
-        schema = CollectionSchema(fields=node_fields,
-                                  description=f"Schema for collection {node_coll_name}")
+
+        schema = self.pymilvus_modules["CollectionSchema"](
+            fields=node_fields, description=f"Schema for collection {node_coll_name}"
+        )
 
         # Create collection if it doesn't exist
-        if not utility.has_collection(node_coll_name):
-            collection = Collection(name=node_coll_name, schema=schema)
+        if not self.pymilvus_modules["utility"].has_collection(node_coll_name):
+            collection = self.pymilvus_modules["Collection"](
+                name=node_coll_name, schema=schema
+            )
         else:
-            collection = Collection(name=node_coll_name)
-
-        # Create indexes
-        collection.create_index(field_name="node_index",
-                                index_params={"index_type": "STL_SORT"},
-                                index_name="node_index_index")
-        collection.create_index(field_name="node_name",
-                                index_params={"index_type": "INVERTED"},
-                                index_name="node_name_index")
-        collection.create_index(field_name="node_type",
-                                index_params={"index_type": "INVERTED"},
-                                index_name="node_type_index")
-        collection.create_index(field_name="desc",
-                                index_params={"index_type": "INVERTED"},
-                                index_name="desc_index")
-        collection.create_index(field_name="desc_emb",
-                                index_params={"index_type": "GPU_CAGRA",
-                                              "metric_type": "IP"},
-                                index_name="desc_emb_index")
+            collection = self.pymilvus_modules["Collection"](name=node_coll_name)
+
+        # Create indexes with dynamic parameters
+        collection.create_index(
+            field_name="node_index",
+            index_params={"index_type": "STL_SORT"},
+            index_name="node_index_index",
+        )
+        collection.create_index(
+            field_name="node_name",
+            index_params={"index_type": "INVERTED"},
+            index_name="node_name_index",
+        )
+        collection.create_index(
+            field_name="node_type",
+            index_params={"index_type": "INVERTED"},
+            index_name="node_type_index",
+        )
+        collection.create_index(
+            field_name="desc",
+            index_params={"index_type": "INVERTED"},
+            index_name="desc_index",
+        )
+        collection.create_index(
+            field_name="desc_emb",
+            index_params={
+                "index_type": self.vector_index_type,
+                "metric_type": self.metric_type,
+            },
+            index_name="desc_emb_index",
+        )
 
         # Prepare and insert data
-        desc_emb_norm = cp.asarray(nodes_df["desc_emb"].list.leaves).astype(cp.float32).\
-            reshape(nodes_df.shape[0], -1)
-        desc_emb_norm = self.normalize_matrix(desc_emb_norm, axis=1)
+        desc_emb_data = self._extract_embeddings(nodes_df, "desc_emb")
+        desc_emb_normalized = self._normalize_matrix(desc_emb_data, axis=1)
+
         data = [
-            nodes_df["node_index"].to_arrow().to_pylist(),
-            nodes_df["node_id"].to_arrow().to_pylist(),
-            nodes_df["node_name"].to_arrow().to_pylist(),
-            nodes_df["node_type"].to_arrow().to_pylist(),
-            nodes_df["desc"].to_arrow().to_pylist(),
-            desc_emb_norm.tolist(), # Use normalized embeddings
+            self._to_list(nodes_df["node_index"]),
+            self._to_list(nodes_df["node_id"]),
+            self._to_list(nodes_df["node_name"]),
+            self._to_list(nodes_df["node_type"]),
+            self._to_list(nodes_df["desc"]),
+            self._to_list(desc_emb_normalized),
         ]
 
         # Insert data in batches
         total = len(data[0])
-        for i in tqdm(range(0, total, self.batch_size), desc="Inserting nodes"):
-            batch = [col[i:i+self.batch_size] for col in data]
+        for i in self.tqdm(range(0, total, self.batch_size), desc="Inserting nodes"):
+            batch = [col[i : i + self.batch_size] for col in data]
             collection.insert(batch)
 
         collection.flush()
-        logger.info("Nodes collection created with %d entities", collection.num_entities)
+        logger.info(
+            "Nodes collection created with %d entities", collection.num_entities
+        )
 
-    def create_node_type_collections(self, nodes_df: cudf.DataFrame):
+    def create_node_type_collections(self, nodes_df):
         """Create separate collections for each node type."""
         logger.info("Creating node type-specific collections...")
 
-        for node_type, nodes_df_ in tqdm(nodes_df.groupby('node_type'),
-                                         desc="Processing node types"):
-            node_coll_name = f"{self.milvus_database}_nodes_{node_type.replace('/', '_')}"
+        for node_type, nodes_df_ in self.tqdm(
+            nodes_df.groupby("node_type"), desc="Processing node types"
+        ):
+            node_coll_name = (
+                f"{self.milvus_database}_nodes_{node_type.replace('/', '_')}"
+            )
+
+            # Get embedding dimensions
+            desc_dim = self._get_embedding_dimension(nodes_df_, "desc_emb")
+            feat_dim = self._get_embedding_dimension(nodes_df_, "feat_emb")
 
             node_fields = [
-                FieldSchema(name="node_index",
-                            dtype=DataType.INT64,
-                            is_primary=True,
-                            auto_id=False),
-                FieldSchema(name="node_id",
-                            dtype=DataType.VARCHAR,
-                            max_length=1024),
-                FieldSchema(name="node_name",
-                            dtype=DataType.VARCHAR,
-                            max_length=1024,
-                            enable_analyzer=True,
-                            enable_match=True),
-                FieldSchema(name="node_type",
-                            dtype=DataType.VARCHAR,
-                            max_length=1024,
-                            enable_analyzer=True,
-                            enable_match=True),
-                FieldSchema(name="desc",
-                            dtype=DataType.VARCHAR,
-                            max_length=40960,
-                            enable_analyzer=True,
-                            enable_match=True),
-                FieldSchema(name="desc_emb",
-                            dtype=DataType.FLOAT_VECTOR,
-                            dim=len(nodes_df_.iloc[0]['desc_emb'].to_arrow().to_pylist()[0])),
-                FieldSchema(name="feat",
-                            dtype=DataType.VARCHAR,
-                            max_length=40960,
-                            enable_analyzer=True,
-                            enable_match=True),
-                FieldSchema(name="feat_emb",
-                            dtype=DataType.FLOAT_VECTOR,
-                            dim=len(nodes_df_.iloc[0]['feat_emb'].to_arrow().to_pylist()[0])),
+                self.pymilvus_modules["FieldSchema"](
+                    name="node_index",
+                    dtype=self.pymilvus_modules["DataType"].INT64,
+                    is_primary=True,
+                    auto_id=False,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="node_id",
+                    dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                    max_length=1024,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="node_name",
+                    dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                    max_length=1024,
+                    enable_analyzer=True,
+                    enable_match=True,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="node_type",
+                    dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                    max_length=1024,
+                    enable_analyzer=True,
+                    enable_match=True,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="desc",
+                    dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                    max_length=40960,
+                    enable_analyzer=True,
+                    enable_match=True,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="desc_emb",
+                    dtype=self.pymilvus_modules["DataType"].FLOAT_VECTOR,
+                    dim=desc_dim,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="feat",
+                    dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                    max_length=40960,
+                    enable_analyzer=True,
+                    enable_match=True,
+                ),
+                self.pymilvus_modules["FieldSchema"](
+                    name="feat_emb",
+                    dtype=self.pymilvus_modules["DataType"].FLOAT_VECTOR,
+                    dim=feat_dim,
+                ),
             ]
-            schema = CollectionSchema(fields=node_fields,
-                                      description=f"schema for collection {node_coll_name}")
 
-            if not utility.has_collection(node_coll_name):
-                collection = Collection(name=node_coll_name, schema=schema)
+            schema = self.pymilvus_modules["CollectionSchema"](
+                fields=node_fields,
+                description=f"schema for collection {node_coll_name}",
+            )
+
+            if not self.pymilvus_modules["utility"].has_collection(node_coll_name):
+                collection = self.pymilvus_modules["Collection"](
+                    name=node_coll_name, schema=schema
+                )
             else:
-                collection = Collection(name=node_coll_name)
-
-            # Create indexes
-            collection.create_index(field_name="node_index",
-                                    index_params={"index_type": "STL_SORT"},
-                                    index_name="node_index_index")
-            collection.create_index(field_name="node_name",
-                                    index_params={"index_type": "INVERTED"},
-                                    index_name="node_name_index")
-            collection.create_index(field_name="node_type",
-                                    index_params={"index_type": "INVERTED"},
-                                    index_name="node_type_index")
-            collection.create_index(field_name="desc",
-                                    index_params={"index_type": "INVERTED"},
-                                    index_name="desc_index")
-            collection.create_index(field_name="desc_emb",
-                                    index_params={"index_type": "GPU_CAGRA",
-                                                  "metric_type": "IP"},
-                                    index_name="desc_emb_index")
-            collection.create_index(field_name="feat_emb",
-                                    index_params={"index_type": "GPU_CAGRA",
-                                                  "metric_type": "IP"},
-                                    index_name="feat_emb_index")
+                collection = self.pymilvus_modules["Collection"](name=node_coll_name)
+
+            # Create indexes with dynamic parameters
+            collection.create_index(
+                field_name="node_index",
+                index_params={"index_type": "STL_SORT"},
+                index_name="node_index_index",
+            )
+            collection.create_index(
+                field_name="node_name",
+                index_params={"index_type": "INVERTED"},
+                index_name="node_name_index",
+            )
+            collection.create_index(
+                field_name="node_type",
+                index_params={"index_type": "INVERTED"},
+                index_name="node_type_index",
+            )
+            collection.create_index(
+                field_name="desc",
+                index_params={"index_type": "INVERTED"},
+                index_name="desc_index",
+            )
+            collection.create_index(
+                field_name="desc_emb",
+                index_params={
+                    "index_type": self.vector_index_type,
+                    "metric_type": self.metric_type,
+                },
+                index_name="desc_emb_index",
+            )
+            collection.create_index(
+                field_name="feat_emb",
+                index_params={
+                    "index_type": self.vector_index_type,
+                    "metric_type": self.metric_type,
+                },
+                index_name="feat_emb_index",
+            )
 
             # Prepare data
-            desc_emb_norm = cp.asarray(nodes_df_["desc_emb"].list.leaves).astype(cp.float32).\
-                reshape(nodes_df_.shape[0], -1)
-            desc_emb_norm = self.normalize_matrix(desc_emb_norm, axis=1)
-            feat_emb_norm = cp.asarray(nodes_df_["feat_emb"].list.leaves).astype(cp.float32).\
-                reshape(nodes_df_.shape[0], -1)
-            feat_emb_norm = self.normalize_matrix(feat_emb_norm, axis=1)
+            desc_emb_data = self._extract_embeddings(nodes_df_, "desc_emb")
+            feat_emb_data = self._extract_embeddings(nodes_df_, "feat_emb")
+
+            desc_emb_normalized = self._normalize_matrix(desc_emb_data, axis=1)
+            feat_emb_normalized = self._normalize_matrix(feat_emb_data, axis=1)
+
             data = [
-                nodes_df_["node_index"].to_arrow().to_pylist(),
-                nodes_df_["node_id"].to_arrow().to_pylist(),
-                nodes_df_["node_name"].to_arrow().to_pylist(),
-                nodes_df_["node_type"].to_arrow().to_pylist(),
-                nodes_df_["desc"].to_arrow().to_pylist(),
-                desc_emb_norm.tolist(), # Use normalized embeddings
-                nodes_df_["feat"].to_arrow().to_pylist(),
-                feat_emb_norm.tolist(), # Use normalized embeddings
+                self._to_list(nodes_df_["node_index"]),
+                self._to_list(nodes_df_["node_id"]),
+                self._to_list(nodes_df_["node_name"]),
+                self._to_list(nodes_df_["node_type"]),
+                self._to_list(nodes_df_["desc"]),
+                self._to_list(desc_emb_normalized),
+                self._to_list(nodes_df_["feat"]),
+                self._to_list(feat_emb_normalized),
             ]
 
             # Insert data in batches
             total_rows = len(data[0])
             for i in range(0, total_rows, self.batch_size):
-                batch = [col[i:i + self.batch_size] for col in data]
+                batch = [col[i : i + self.batch_size] for col in data]
                 collection.insert(batch)
 
             collection.flush()
-            logger.info("Collection %s created with %d entities",
-                        node_coll_name, collection.num_entities)
+            logger.info(
+                "Collection %s created with %d entities",
+                node_coll_name,
+                collection.num_entities,
+            )
 
-    def create_edges_collection(self,
-                                edges_enrichment_df: cudf.DataFrame,
-                                edges_embedding_df: List[cudf.DataFrame]):
-        """Create and populate the edges collection."""
+    def create_edges_collection(self, edges_enrichment_df, edges_embedding_df: List):
+        """Create and populate the edges collection - exact original logic."""
         logger.info("Creating edges collection...")
 
         edge_coll_name = f"{self.milvus_database}_edges"
 
+        # Get embedding dimension from first chunk - exact original logic
+        if self.use_gpu:
+            emb_dim = len(
+                edges_embedding_df[0].loc[0, "edge_emb"]
+            )  # Original cudf access
+        else:
+            first_edge_emb = edges_embedding_df[0].iloc[0]["edge_emb"]
+            emb_dim = (
+                len(first_edge_emb)
+                if isinstance(first_edge_emb, list)
+                else len(first_edge_emb.tolist())
+            )
+
         edge_fields = [
-            FieldSchema(name="triplet_index",
-                        dtype=DataType.INT64,
-                        is_primary=True,
-                        auto_id=False),
-            FieldSchema(name="head_id",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024),
-            FieldSchema(name="head_index",
-                        dtype=DataType.INT64),
-            FieldSchema(name="tail_id",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024),
-            FieldSchema(name="tail_index",
-                        dtype=DataType.INT64),
-            FieldSchema(name="edge_type",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024),
-            FieldSchema(name="display_relation",
-                        dtype=DataType.VARCHAR,
-                        max_length=1024),
-            FieldSchema(name="feat",
-                        dtype=DataType.VARCHAR,
-                        max_length=40960),
-            FieldSchema(name="feat_emb",
-                        dtype=DataType.FLOAT_VECTOR,
-                        dim=len(edges_embedding_df[0].loc[0, 'edge_emb'])),
+            self.pymilvus_modules["FieldSchema"](
+                name="triplet_index",
+                dtype=self.pymilvus_modules["DataType"].INT64,
+                is_primary=True,
+                auto_id=False,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="head_id",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="head_index", dtype=self.pymilvus_modules["DataType"].INT64
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="tail_id",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="tail_index", dtype=self.pymilvus_modules["DataType"].INT64
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="edge_type",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="display_relation",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=1024,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="feat",
+                dtype=self.pymilvus_modules["DataType"].VARCHAR,
+                max_length=40960,
+            ),
+            self.pymilvus_modules["FieldSchema"](
+                name="feat_emb",
+                dtype=self.pymilvus_modules["DataType"].FLOAT_VECTOR,
+                dim=emb_dim,
+            ),
         ]
-        edge_schema = CollectionSchema(fields=edge_fields,
-                                       description="Schema for edges collection")
 
-        if not utility.has_collection(edge_coll_name):
-            collection = Collection(name=edge_coll_name, schema=edge_schema)
+        edge_schema = self.pymilvus_modules["CollectionSchema"](
+            fields=edge_fields, description="Schema for edges collection"
+        )
+
+        if not self.pymilvus_modules["utility"].has_collection(edge_coll_name):
+            collection = self.pymilvus_modules["Collection"](
+                name=edge_coll_name, schema=edge_schema
+            )
         else:
-            collection = Collection(name=edge_coll_name)
-
-        # Create indexes
-        collection.create_index(field_name="triplet_index",
-                                index_params={"index_type": "STL_SORT"},
-                                index_name="triplet_index_index")
-        collection.create_index(field_name="head_index",
-                                index_params={"index_type": "STL_SORT"},
-                                index_name="head_index_index")
-        collection.create_index(field_name="tail_index",
-                                index_params={"index_type": "STL_SORT"},
-                                index_name="tail_index_index")
-        collection.create_index(field_name="feat_emb",
-                                index_params={"index_type": "GPU_CAGRA",
-                                              "metric_type": "IP"},
-                                index_name="feat_emb_index")
-
-        # Iterate over chunked edges embedding df
-        for edges_df in tqdm(edges_embedding_df):
+            collection = self.pymilvus_modules["Collection"](name=edge_coll_name)
+
+        # Create indexes with dynamic parameters
+        collection.create_index(
+            field_name="triplet_index",
+            index_params={"index_type": "STL_SORT"},
+            index_name="triplet_index_index",
+        )
+        collection.create_index(
+            field_name="head_index",
+            index_params={"index_type": "STL_SORT"},
+            index_name="head_index_index",
+        )
+        collection.create_index(
+            field_name="tail_index",
+            index_params={"index_type": "STL_SORT"},
+            index_name="tail_index_index",
+        )
+        collection.create_index(
+            field_name="feat_emb",
+            index_params={
+                "index_type": self.vector_index_type,
+                "metric_type": self.metric_type,
+            },
+            index_name="feat_emb_index",
+        )
+
+        # Iterate over chunked edges embedding df - exact original logic
+        for edges_df in self.tqdm(edges_embedding_df, desc="Processing edge chunks"):
             # Merge enrichment with embedding
             merged_edges_df = edges_enrichment_df.merge(
-                edges_df[["triplet_index", "edge_emb"]],
-                on="triplet_index",
-                how="inner"
+                edges_df[["triplet_index", "edge_emb"]], on="triplet_index", how="inner"
             )
 
-            # Prepare data
-            edge_emb_cp = cp.asarray(merged_edges_df["edge_emb"].list.leaves).astype(cp.float32).\
-                reshape(merged_edges_df.shape[0], -1)
-            edge_emb_norm = self.normalize_matrix(edge_emb_cp, axis=1)
+            # Prepare embeddings - exact original logic for GPU
+            if self.use_gpu:
+                edge_emb_cp = (
+                    self.cp.asarray(merged_edges_df["edge_emb"].list.leaves)
+                    .astype(self.cp.float32)
+                    .reshape(merged_edges_df.shape[0], -1)
+                )
+                edge_emb_norm = self._normalize_matrix(edge_emb_cp, axis=1)
+            else:
+                edge_emb_data = self._extract_embeddings(merged_edges_df, "edge_emb")
+                edge_emb_norm = self._normalize_matrix(edge_emb_data, axis=1)
+
             data = [
-                merged_edges_df["triplet_index"].to_arrow().to_pylist(),
-                merged_edges_df["head_id"].to_arrow().to_pylist(),
-                merged_edges_df["head_index"].to_arrow().to_pylist(),
-                merged_edges_df["tail_id"].to_arrow().to_pylist(),
-                merged_edges_df["tail_index"].to_arrow().to_pylist(),
-                merged_edges_df["edge_type_str"].to_arrow().to_pylist(),
-                merged_edges_df["display_relation"].to_arrow().to_pylist(),
-                merged_edges_df["feat"].to_arrow().to_pylist(),
-                edge_emb_norm.tolist(), # Use normalized embeddings
+                self._to_list(merged_edges_df["triplet_index"]),
+                self._to_list(merged_edges_df["head_id"]),
+                self._to_list(merged_edges_df["head_index"]),
+                self._to_list(merged_edges_df["tail_id"]),
+                self._to_list(merged_edges_df["tail_index"]),
+                self._to_list(merged_edges_df["edge_type_str"]),  # Original field name
+                self._to_list(merged_edges_df["display_relation"]),
+                self._to_list(merged_edges_df["feat"]),
+                self._to_list(edge_emb_norm),
             ]
 
             # Insert data in batches
             total = len(data[0])
-            for i in tqdm(range(0, total, self.batch_size), desc="Inserting edges"):
-                batch_data = [d[i:i+self.batch_size] for d in data]
+            for i in self.tqdm(
+                range(0, total, self.batch_size), desc="Inserting edges"
+            ):
+                batch_data = [d[i : i + self.batch_size] for d in data]
                 collection.insert(batch_data)
 
         collection.flush()
-        logger.info("Edges collection created with %d entities", collection.num_entities)
+        logger.info(
+            "Edges collection created with %d entities", collection.num_entities
+        )
 
     def run(self):
         """Main execution method."""
         try:
-            logger.info("Starting Milvus data loading process...")
+            logger.info("Starting Dynamic Milvus data loading process...")
+            logger.info(
+                "System: %s %s", self.detector.os_type, self.detector.architecture
+            )
+            logger.info("GPU acceleration: %s", self.use_gpu)
 
             # Connect to Milvus
             self.connect_to_milvus()
@@ -443,66 +801,110 @@ class MilvusDataLoader:
             # Prepare data
             logger.info("Data Preparation started...")
             # Get nodes enrichment and embedding dataframes
-            nodes_enrichment_df = graph['nodes']['enrichment']
-            nodes_embedding_df = graph['nodes']['embedding']
+            nodes_enrichment_df = graph["nodes"]["enrichment"]
+            nodes_embedding_df = graph["nodes"]["embedding"]
 
             # Get edges enrichment and embedding dataframes
-            edges_enrichment_df = graph['edges']['enrichment']
-            # !!consisted of a list of dataframes!!
-            edges_embedding_df = graph['edges']['embedding']
+            edges_enrichment_df = graph["edges"]["enrichment"]
+            edges_embedding_df = graph["edges"]["embedding"]  # List of dataframes
 
-            # For nodes, we can directly merge enrichment and embedding
             # Merge nodes enrichment and embedding dataframes
             merged_nodes_df = nodes_enrichment_df.merge(
                 nodes_embedding_df[["node_id", "desc_emb", "feat_emb"]],
                 on="node_id",
-                how="left"
+                how="left",
             )
 
             # Create collections and load data
             self.create_nodes_collection(merged_nodes_df)
             self.create_node_type_collections(merged_nodes_df)
-            self.create_edges_collection(edges_enrichment_df,
-                                         edges_embedding_df)
+            self.create_edges_collection(edges_enrichment_df, edges_embedding_df)
 
             # List all collections for verification
             logger.info("Data loading completed successfully!")
             logger.info("Created collections:")
-            for coll in utility.list_collections():
-                collection = Collection(name=coll)
+            for coll in self.pymilvus_modules["utility"].list_collections():
+                collection = self.pymilvus_modules["Collection"](name=coll)
                 logger.info("  %s: %d entities", coll, collection.num_entities)
 
         except Exception as e:
             logger.error("Error during data loading: %s", str(e))
+            import traceback
+
+            logger.error("Full traceback: %s", traceback.format_exc())
             raise
 
 
 def main():
-    """Main function to run the data loader."""
+    """Main function to run the dynamic data loader."""
     # Resolve the fallback data path relative to this script's location
     script_dir = os.path.dirname(os.path.abspath(__file__))
     default_data_dir = os.path.join(script_dir, "tests/files/biobridge_multimodal/")
 
-    # Configuration
+    # Configuration with environment variable fallbacks - matches original exactly
     config = {
-        'milvus_host': os.getenv('MILVUS_HOST', 'localhost'),
-        'milvus_port': os.getenv('MILVUS_PORT', '19530'),
-        'milvus_user': os.getenv('MILVUS_USER', 'root'),
-        'milvus_password': os.getenv('MILVUS_PASSWORD', 'Milvus'),
-        'milvus_database': os.getenv('MILVUS_DATABASE', 't2kg_primekg'),
-        'data_dir': os.getenv('DATA_DIR', default_data_dir),
-        'batch_size': int(os.getenv('BATCH_SIZE', '500')),
-        'chunk_size': int(os.getenv('CHUNK_SIZE', '5')),
+        "milvus_host": os.getenv("MILVUS_HOST", "localhost"),
+        "milvus_port": os.getenv("MILVUS_PORT", "19530"),
+        "milvus_user": os.getenv("MILVUS_USER", "root"),
+        "milvus_password": os.getenv("MILVUS_PASSWORD", "Milvus"),
+        "milvus_database": os.getenv("MILVUS_DATABASE", "t2kg_primekg"),
+        "data_dir": os.getenv("DATA_DIR", default_data_dir),
+        "batch_size": int(os.getenv("BATCH_SIZE", "500")),
+        "chunk_size": int(os.getenv("CHUNK_SIZE", "5")),
+        "auto_install_packages": os.getenv("AUTO_INSTALL_PACKAGES", "true").lower()
+        == "true",
     }
 
-    # Print configuration for debugging
-    print("[DATA LOADER] Configuration:")
-    for key, value in config.items():
-        print(f"[DATA LOADER]   {key}: {value}")
+    # Override detection for testing/forcing specific modes
+    force_cpu = os.getenv("FORCE_CPU", "false").lower() == "true"
+    if force_cpu:
+        logger.info("FORCE_CPU environment variable set - forcing CPU mode")
 
-    # Create and run data loader
-    loader = MilvusDataLoader(config)
-    loader.run()
+    # Print configuration for debugging - matches original format
+    logger.info("=== Dynamic Milvus Data Loader ===")
+    logger.info("Configuration:")
+    for key, value in config.items():
+        # Don't log sensitive information
+        if "password" in key.lower():
+            logger.info("  %s: %s", key, "*" * len(str(value)))
+        else:
+            logger.info("  %s: %s", key, value)
+
+    # Additional environment info
+    logger.info("Environment:")
+    logger.info("  Python version: %s", sys.version)
+    logger.info("  Platform: %s", platform.platform())
+    logger.info("  Force CPU mode: %s", force_cpu)
+    logger.info("  Script directory: %s", script_dir)
+    logger.info("  Default data directory: %s", default_data_dir)
+
+    try:
+        # Create and run dynamic data loader
+        loader = DynamicDataLoader(config)
+
+        # Override GPU detection if forced
+        if force_cpu:
+            loader.detector.use_gpu = False
+            loader.use_gpu = False
+            loader.normalize_vectors = False
+            loader.vector_index_type = "HNSW"
+            loader.metric_type = "COSINE"
+            logger.info("Forced CPU mode - updated loader settings")
+
+        # Run the data loading process
+        loader.run()
+
+        logger.info("=== Data Loading Completed Successfully ===")
+
+    except KeyboardInterrupt:
+        logger.info("Data loading interrupted by user")
+        sys.exit(1)
+    except Exception as e:
+        logger.error("Fatal error during data loading: %s", str(e))
+        import traceback
+
+        logger.error("Full traceback: %s", traceback.format_exc())
+        sys.exit(1)
 
 
 if __name__ == "__main__":