mirror of
https://github.com/microsoft/graphrag.git
synced 2025-07-04 07:26:30 +00:00
47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
|
# Copyright (c) 2024 Microsoft Corporation.
|
||
|
|
# Licensed under the MIT License
|
||
|
|
import asyncio
|
||
|
|
import os
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
|
||
|
|
from graphrag.index 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())
|