Spaces:
Sleeping
Sleeping
File size: 634 Bytes
70ab3b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | from nano_graphrag import GraphRAG, QueryParam
graph_func = GraphRAG(
working_dir="../bedrock_example",
using_amazon_bedrock=True,
best_model_id="us.anthropic.claude-3-sonnet-20240229-v1:0",
cheap_model_id="us.anthropic.claude-3-haiku-20240307-v1:0",
)
with open("../tests/mock_data.txt") as f:
graph_func.insert(f.read())
prompt = "What are the top themes in this story?"
# Perform global graphrag search
print(graph_func.query(prompt, param=QueryParam(mode="global")))
# Perform local graphrag search (I think is better and more scalable one)
print(graph_func.query(prompt, param=QueryParam(mode="local")))
|