graphrag/examples/custom_input/run.py

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
import asyncio
import os

import pandas as pd

from graphrag.index.run import run_pipeline_with_config

pipeline_file = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "./pipeline.yml"
)


async def run():
    # Load your dataset
    dataset = _load_dataset_some_unique_way()

    # Load your config without the input section
    config = pipeline_file

    # Grab the last result from the pipeline, should be our entity extraction
    outputs = []
    async for output in run_pipeline_with_config(
        config_or_path=config, dataset=dataset
    ):
        outputs.append(output)
    pipeline_result = outputs[-1]

    if pipeline_result.result is not None:
        # Should look something like
        #            col1  col2 filled_column
        # 0     2     4  Filled Value
        # 1     5    10  Filled Value
        print(pipeline_result.result)
    else:
        print("No results!")


def _load_dataset_some_unique_way() -> pd.DataFrame:
    # Totally loaded from some other place
    return pd.DataFrame([{"col1": 2, "col2": 4}, {"col1": 5, "col2": 10}])


if __name__ == "__main__":
    asyncio.run(run())