check_entities.py

# /home/wangxinxin/my_graphrag/check_entities.py
import pandas as pd
import os

FILE_PATH = "output/entities.parquet"

if os.path.exists(FILE_PATH):
    df = pd.read_parquet(FILE_PATH)
    
    # Search for entity names
    keywords = ["AHG", "顶峰", "量子", "星云"]
    print(f"Searching entity database, keywords: {keywords} ...")
    print("-" * 50)
    
    found = False
    for index, row in df.iterrows():
        name = row['title']
        # Check if the name contains any keywords
        if any(k in name for k in keywords):
            print(f"Found entity: [{name}] (ID: {row.get('human_readable_id', 'N/A')})")
            print(f"Description: {row.get('description', '')[:50]}...")
            found = True
            
    if not found:
        print("No relevant entities found! The LLM might have missed them or categorized them as non-entities.")
else:
    print(f"File does not exist: {FILE_PATH}")