fix demo

examples/vram_management_demo.py

@@ -1,8 +1,10 @@
 import os
 import time
+import asyncio
 from lightrag import LightRAG, QueryParam
 from lightrag.llm.ollama import ollama_model_complete, ollama_embed
 from lightrag.utils import EmbeddingFunc
+from lightrag.kg.shared_storage import initialize_pipeline_status
 
 # Working directory and the directory path for text files
 WORKING_DIR = "./dickens"
@@ -12,17 +14,22 @@ TEXT_FILES_DIR = "/llm/mt"
 if not os.path.exists(WORKING_DIR):
     os.mkdir(WORKING_DIR)
 
-# Initialize LightRAG
-rag = LightRAG(
-    working_dir=WORKING_DIR,
-    llm_model_func=ollama_model_complete,
-    llm_model_name="qwen2.5:3b-instruct-max-context",
-    embedding_func=EmbeddingFunc(
-        embedding_dim=768,
-        max_token_size=8192,
-        func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text"),
-    ),
-)
+async def initialize_rag():
+    # Initialize LightRAG
+    rag = LightRAG(
+        working_dir=WORKING_DIR,
+        llm_model_func=ollama_model_complete,
+        llm_model_name="qwen2.5:3b-instruct-max-context",
+        embedding_func=EmbeddingFunc(
+            embedding_dim=768,
+            max_token_size=8192,
+            func=lambda texts: ollama_embed(texts, embed_model="nomic-embed-text"),
+        ),
+    )
+    await rag.initialize_storages()
+    await initialize_pipeline_status()
+
+    return rag
 
 # Read all .txt files from the TEXT_FILES_DIR directory
 texts = []
@@ -47,58 +54,65 @@ def insert_texts_with_retry(rag, texts, retries=3, delay=5):
     raise RuntimeError("Failed to insert texts after multiple retries.")
 
 
-insert_texts_with_retry(rag, texts)
+def main():
+    # Initialize RAG instance
+    rag = asyncio.run(initialize_rag())
 
-# Perform different types of queries and handle potential errors
-try:
-    print(
-        rag.query(
-            "What are the top themes in this story?", param=QueryParam(mode="naive")
+    insert_texts_with_retry(rag, texts)
+
+    # Perform different types of queries and handle potential errors
+    try:
+        print(
+            rag.query(
+                "What are the top themes in this story?", param=QueryParam(mode="naive")
+            )
         )
-    )
-except Exception as e:
-    print(f"Error performing naive search: {e}")
+    except Exception as e:
+        print(f"Error performing naive search: {e}")
 
-try:
-    print(
-        rag.query(
-            "What are the top themes in this story?", param=QueryParam(mode="local")
+    try:
+        print(
+            rag.query(
+                "What are the top themes in this story?", param=QueryParam(mode="local")
+            )
         )
-    )
-except Exception as e:
-    print(f"Error performing local search: {e}")
+    except Exception as e:
+        print(f"Error performing local search: {e}")
 
-try:
-    print(
-        rag.query(
-            "What are the top themes in this story?", param=QueryParam(mode="global")
+    try:
+        print(
+            rag.query(
+                "What are the top themes in this story?", param=QueryParam(mode="global")
+            )
         )
-    )
-except Exception as e:
-    print(f"Error performing global search: {e}")
+    except Exception as e:
+        print(f"Error performing global search: {e}")
 
-try:
-    print(
-        rag.query(
-            "What are the top themes in this story?", param=QueryParam(mode="hybrid")
+    try:
+        print(
+            rag.query(
+                "What are the top themes in this story?", param=QueryParam(mode="hybrid")
+            )
         )
-    )
-except Exception as e:
-    print(f"Error performing hybrid search: {e}")
+    except Exception as e:
+        print(f"Error performing hybrid search: {e}")
 
 
-# Function to clear VRAM resources
-def clear_vram():
-    os.system("sudo nvidia-smi --gpu-reset")
+    # Function to clear VRAM resources
+    def clear_vram():
+        os.system("sudo nvidia-smi --gpu-reset")
 
 
-# Regularly clear VRAM to prevent overflow
-clear_vram_interval = 3600  # Clear once every hour
-start_time = time.time()
+    # Regularly clear VRAM to prevent overflow
+    clear_vram_interval = 3600  # Clear once every hour
+    start_time = time.time()
 
-while True:
-    current_time = time.time()
-    if current_time - start_time > clear_vram_interval:
-        clear_vram()
-        start_time = current_time
-    time.sleep(60)  # Check the time every minute
+    while True:
+        current_time = time.time()
+        if current_time - start_time > clear_vram_interval:
+            clear_vram()
+            start_time = current_time
+        time.sleep(60)  # Check the time every minute
+
+if __name__ == "__main__":
+    main()