ContextQMD
Back to Graphrag

Copyright (c) 2024 Microsoft Corporation.

docs/examples_notebooks/input_documents.ipynb

3.0.93.0 KB
Original Source
python
# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License.

Example of indexing from an existing in-memory dataframe

Newer versions of GraphRAG let you submit a dataframe directly instead of running through the input processing step. This notebook demonstrates with regular or update runs.

If performing an update, the assumption is that your dataframe contains only the new documents to add to the index.

python
from pathlib import Path
from pprint import pprint

import graphrag.api as api
import pandas as pd
from graphrag.config.load_config import load_config
from graphrag.index.typing.pipeline_run_result import PipelineRunResult
python
PROJECT_DIRECTORY = "<your project directory>"
UPDATE = False
FILENAME = "new_documents.parquet" if UPDATE else "<original_documents>.parquet"
inputs = pd.read_parquet(f"{PROJECT_DIRECTORY}/input/{FILENAME}")
# Only the bare minimum for input. These are the same fields that would be present after the load_input_documents workflow
inputs = inputs.loc[:, ["id", "title", "text", "creation_date"]]

Generate a GraphRagConfig object

python
graphrag_config = load_config(Path(PROJECT_DIRECTORY))

Indexing API

Indexing is the process of ingesting raw text data and constructing a knowledge graph. GraphRAG currently supports plaintext (.txt) and .csv file formats.

Build an index

python
index_result: list[PipelineRunResult] = await api.build_index(
    config=graphrag_config, input_documents=inputs, is_update_run=UPDATE
)

# index_result is a list of workflows that make up the indexing pipeline that was run
for workflow_result in index_result:
    status = f"error\n{workflow_result.errors}" if workflow_result.errors else "success"
    print(f"Workflow Name: {workflow_result.workflow}\tStatus: {status}")

Query an index

To query an index, several index files must first be read into memory and passed to the query API.

python
entities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/entities.parquet")
communities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/communities.parquet")
community_reports = pd.read_parquet(
    f"{PROJECT_DIRECTORY}/output/community_reports.parquet"
)

response, context = await api.global_search(
    config=graphrag_config,
    entities=entities,
    communities=communities,
    community_reports=community_reports,
    community_level=2,
    dynamic_community_selection=False,
    response_type="Multiple Paragraphs",
    query="What are the top five themes of the dataset?",
)

The response object is the official reponse from graphrag while the context object holds various metadata regarding the querying process used to obtain the final response.

python
print(response)

Digging into the context a bit more provides users with extremely granular information such as what sources of data (down to the level of text chunks) were ultimately retrieved and used as part of the context sent to the LLM model).

python
pprint(context)  # noqa: T203