File size: 2,319 Bytes
af45684 1f844c6 af45684 cb56de7 af45684 1f844c6 af45684 cb56de7 af45684 1f844c6 cb56de7 af45684 1f844c6 cb56de7 af45684 1f844c6 cb56de7 af45684 1f844c6 cb56de7 1f844c6 cb56de7 1f844c6 cb56de7 1f844c6 af45684 cb56de7 af45684 cb56de7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import gpt_4o_mini_complete, gpt_4o_complete, openai_embed
from lightrag.kg.shared_storage import initialize_pipeline_status
WORKING_DIR = "./lightrag_demo"
if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
async def initialize_rag():
rag = LightRAG(
working_dir=WORKING_DIR,
embedding_func=openai_embed,
llm_model_func=gpt_4o_mini_complete, # Default model for queries
)
await rag.initialize_storages()
await initialize_pipeline_status()
return rag
def main():
# Initialize RAG instance
rag = asyncio.run(initialize_rag())
# Load the data
with open("./book.txt", "r", encoding="utf-8") as f:
rag.insert(f.read())
# Query with naive mode (default model)
print("--- NAIVE mode ---")
print(
rag.query(
"What are the main themes in this story?", param=QueryParam(mode="naive")
)
)
# Query with local mode (default model)
print("\n--- LOCAL mode ---")
print(
rag.query(
"What are the main themes in this story?", param=QueryParam(mode="local")
)
)
# Query with global mode (default model)
print("\n--- GLOBAL mode ---")
print(
rag.query(
"What are the main themes in this story?", param=QueryParam(mode="global")
)
)
# Query with hybrid mode (default model)
print("\n--- HYBRID mode ---")
print(
rag.query(
"What are the main themes in this story?", param=QueryParam(mode="hybrid")
)
)
# Query with mix mode (default model)
print("\n--- MIX mode ---")
print(
rag.query(
"What are the main themes in this story?", param=QueryParam(mode="mix")
)
)
# Query with a custom model (gpt-4o) for a more complex question
print("\n--- Using custom model for complex analysis ---")
print(
rag.query(
"How does the character development reflect Victorian-era attitudes?",
param=QueryParam(
mode="global",
model_func=gpt_4o_complete, # Override default model with more capable one
),
)
)
if __name__ == "__main__":
main()
|