Update graphrag_agent.py

graphrag_agent.py

@@ -700,35 +700,44 @@ def whether_to_interact(state):
 
 # 数据存放路径
 DATA_DIR = "data"
-ZIP_FILE = "data.zip"
-
-# ✅ 如果 data/ 已存在且非空，就跳过下载
-if os.path.exists(DATA_DIR) and any(os.scandir(DATA_DIR)):
-    print("✅ 已检测到本地 data 文件夹，跳过下载。")
-
-else:
-    # 如果没有解压好的 data，但有 data.zip，则直接解压
-    if os.path.exists(ZIP_FILE):
-        print("📦 检测到本地 data.zip，正在解压...")
-        with ZipFile(ZIP_FILE, "r") as zip_ref:
-            zip_ref.extractall(DATA_DIR)
-        print("✅ 已成功解压本地 data.zip")
+os.makedirs(DATA_DIR, exist_ok=True)
+
+# Hugging Face Dataset repo ID
+REPO_ID = "achenyx1412/DGADIS"
+
+# 需要下载的文件列表
+FILES = [
+    "faiss_node+desc.index",
+    "faiss_node+desc.pkl",
+    "faiss_node.index",
+    "faiss_node.pkl",
+    "faiss_triple3.index",
+    "faiss_triple3.pkl",
+    "kg.gpickle"
+]
 
-    else:
-        # 如果连 data.zip 都没有，才从 HF 下载
-        print("🌐 未检测到本地数据，开始从 Hugging Face 下载 data.zip...")
-        zip_path = hf_hub_download(
-            repo_id="achenyx1412/DGADIS",
-            filename="data.zip",
+# 遍历文件，逐个下载
+for file_name in FILES:
+    local_path = os.path.join(DATA_DIR, file_name)
+    
+    # 如果本地已存在，则跳过下载
+    if os.path.exists(local_path):
+        print(f"✅ 已检测到本地文件 {file_name}，跳过下载。")
+        continue
+    
+    print(f"🌐 正在从 Hugging Face 下载 {file_name} ...")
+    try:
+        hf_hub_download(
+            repo_id=REPO_ID,
+            filename=file_name,
             repo_type="dataset",
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            local_dir=DATA_DIR,
+            local_dir_use_symlinks=False  # 防止 symlink 问题
         )
-
-        # 解压
-        print("📦 正在解压 data.zip...")
-        with ZipFile(zip_path, "r") as zip_ref:
-            zip_ref.extractall(DATA_DIR)
-        print("✅ 已成功下载并解压 data.zip")
+        print(f"✅ 已成功下载 {file_name}")
+    except Exception as e:
+        print(f"❌ 下载 {file_name} 失败: {e}")
 
 def neo4j_retrieval(state: MyState):
     logger.info("---NODE: neo4j_retrieval---")