From f87c235a4c00d13fac047a848cfb70e2346e1524 Mon Sep 17 00:00:00 2001 From: omdivyatej Date: Sun, 23 Mar 2025 21:42:56 +0530 Subject: [PATCH] less comments --- .../lightrag_multi_model_all_modes_demo.py | 111 +++++++++--------- 1 file changed, 54 insertions(+), 57 deletions(-) diff --git a/examples/lightrag_multi_model_all_modes_demo.py b/examples/lightrag_multi_model_all_modes_demo.py index 04adf642..c2f9c3d2 100644 --- a/examples/lightrag_multi_model_all_modes_demo.py +++ b/examples/lightrag_multi_model_all_modes_demo.py @@ -3,22 +3,17 @@ 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 -from lightrag.utils import setup_logger -setup_logger("lightrag", level="INFO") - -WORKING_DIR = "./all_modes_demo" +WORKING_DIR = "./lightrag_demo" if not os.path.exists(WORKING_DIR): os.mkdir(WORKING_DIR) - async def initialize_rag(): - # Initialize LightRAG with a base model (gpt-4o-mini) rag = LightRAG( working_dir=WORKING_DIR, embedding_func=openai_embed, - llm_model_func=gpt_4o_mini_complete, # Default model for most queries + llm_model_func=gpt_4o_mini_complete, # Default model for queries ) await rag.initialize_storages() @@ -26,7 +21,6 @@ async def initialize_rag(): return rag - def main(): # Initialize RAG instance rag = asyncio.run(initialize_rag()) @@ -34,60 +28,63 @@ def main(): # Load the data with open("./book.txt", "r", encoding="utf-8") as f: rag.insert(f.read()) - - # Example query - query_text = "What are the main themes in this story?" - - # Demonstrate using default model (gpt-4o-mini) for all modes - print("\n===== Default Model (gpt-4o-mini) =====") - - for mode in ["local", "global", "hybrid", "naive", "mix"]: - print(f"\n--- {mode.upper()} mode with default model ---") - response = rag.query( - query_text, - param=QueryParam(mode=mode) + + # Query with naive mode (default model) + print("--- NAIVE mode ---") + print( + rag.query( + "What are the main themes in this story?", + param=QueryParam(mode="naive") ) - print(response) - - # Demonstrate using custom model (gpt-4o) for all modes - print("\n===== Custom Model (gpt-4o) =====") - - for mode in ["local", "global", "hybrid", "naive", "mix"]: - print(f"\n--- {mode.upper()} mode with custom model ---") - response = rag.query( - query_text, + ) + + # 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=mode, - model_func=gpt_4o_complete # Override with more capable model + mode="global", + model_func=gpt_4o_complete # Override default model with more capable one ) ) - print(response) - - # Mixed approach - use different models for different modes - print("\n===== Strategic Model Selection =====") - - # Complex analytical question - complex_query = "How does the character development in the story reflect Victorian-era social values?" - - # Use default model for simpler modes - print("\n--- NAIVE mode with default model (suitable for simple retrieval) ---") - response1 = rag.query( - complex_query, - param=QueryParam(mode="naive") # Use default model for basic retrieval ) - print(response1) - - # Use more capable model for complex modes - print("\n--- HYBRID mode with more capable model (for complex analysis) ---") - response2 = rag.query( - complex_query, - param=QueryParam( - mode="hybrid", - model_func=gpt_4o_complete # Use more capable model for complex analysis - ) - ) - print(response2) - if __name__ == "__main__": main() \ No newline at end of file