fix: fix lint problem

Author: ChenZiHong-Gavin

graphgen/operators/__init__.py

@@ -1,21 +1,21 @@
 from .build_kg import BuildKGService
 from .chunk import ChunkService
-from .extract import extract_info
+from .extract import extract
 from .generate import generate_qas
-from .partition import partition_kg
+from .judge import JudgeService
+from .partition import PartitionService
 from .quiz import QuizService
 from .read import read
 from .search import search_all
-from .judge import JudgeService
 
 operators = {
     "read": read,
     "chunk": ChunkService,
     "build_kg": BuildKGService,
     "quiz": QuizService,
     "judge": JudgeService,
-    "extract_info": extract_info,
+    "extract_info": extract,
     "search_all": search_all,
-    "partition_kg": partition_kg,
+    "partition": PartitionService,
     "generate_qas": generate_qas,
 }

graphgen/operators/chunk/chunk_service.py

@@ -94,8 +94,6 @@ def chunk_documents(self, new_docs: list) -> list:
                         **doc,
                     }
                 )
-        self.chunk_storage.upsert(
-            {chunk["_chunk_id"]: chunk for chunk in chunks}
-        )
+        self.chunk_storage.upsert({chunk["_chunk_id"]: chunk for chunk in chunks})
         self.chunk_storage.index_done_callback()
         return chunks

graphgen/operators/extract/__init__.py

@@ -1 +1 @@
-from .extract import extract_info
+from .extract import extract

graphgen/operators/extract/extract.py

@@ -7,7 +7,7 @@
 from graphgen.utils import logger, run_concurrent
 
 
-async def extract_info(
+async def extract(
     llm_client: BaseLLMWrapper,
     chunk_storage: BaseKVStorage,
     extract_config: dict,

graphgen/operators/generate/__init__.py

@@ -1 +1 @@
-from .generate_qas import generate_qas
+from .generate import generate_qas

graphgen/operators/generate/generate.py

@@ -52,7 +52,6 @@ async def generate_qas(
         batches,
         desc="[4/4]Generating QAs",
         unit="batch",
-        progress_bar=progress_bar,
     )
 
     # format

graphgen/operators/judge/judge_service.py

@@ -59,13 +59,11 @@ def judge(self, items: list[dict]) -> None:
             if isinstance(index, str):
                 node_id = index
                 node_data = self.graph_storage.get_node(node_id)
-                if node_data:
-                    node_data["loss"] = loss
-                    self.graph_storage.update_node(node_id, node_data)
+                node_data["loss"] = loss
+                self.graph_storage.update_node(node_id, node_data)
             elif isinstance(index, tuple):
                 edge_source, edge_target = index
                 edge_data = self.graph_storage.get_edge(edge_source, edge_target)
-                if edge_data:
-                    edge_data["loss"] = loss
-                    self.graph_storage.update_edge(edge_source, edge_target, edge_data)
+                edge_data["loss"] = loss
+                self.graph_storage.update_edge(edge_source, edge_target, edge_data)
         self.graph_storage.index_done_callback()

graphgen/operators/partition/partition_service.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Iterable
+from typing import Iterable
 
 import pandas as pd
 
@@ -101,7 +101,7 @@ def _pre_tokenize(self) -> None:
                     node_data["length"] = len(tokens)
                     self.kg_instance.update_node(node_id, node_data)
                 except Exception as e:
-                    logger.warning(f"Failed to tokenize node {node_id}: {e}")
+                    logger.warning("Failed to tokenize node %s: %s", node_id, e)
                     node_data["length"] = 0
 
         # Process edges
@@ -113,7 +113,7 @@ def _pre_tokenize(self) -> None:
                     edge_data["length"] = len(tokens)
                     self.kg_instance.update_edge(u, v, edge_data)
                 except Exception as e:
-                    logger.warning(f"Failed to tokenize edge {u}-{v}: {e}")
+                    logger.warning("Failed to tokenize edge %s-%s: %s", u, v, e)
                     edge_data["length"] = 0
 
         # Persist changes

graphgen/operators/quiz/quiz_service.py

@@ -5,7 +5,7 @@
 from graphgen.bases import BaseGraphStorage, BaseKVStorage, BaseLLMWrapper
 from graphgen.common import init_llm, init_storage
 from graphgen.models import QuizGenerator
-from graphgen.utils import compute_content_hash, logger, run_concurrent
+from graphgen.utils import compute_dict_hash, logger, run_concurrent
 
 
 class QuizService:
@@ -20,7 +20,7 @@ def __init__(
         self.graph_storage: BaseGraphStorage = init_storage(
             backend="networkx", working_dir=working_dir, namespace="graph"
         )
-        # { _description_id: { "description": str, "quizzes": List[Tuple[str, str]] } }
+        # { _quiz_id: { "description": str, "quizzes": List[Tuple[str, str]] } }
         self.quiz_storage: BaseKVStorage = init_storage(
             backend="json_kv", working_dir=working_dir, namespace="quiz"
         )
@@ -37,8 +37,8 @@ def __call__(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
     async def _process_single_quiz(self, item: tuple) -> dict | None:
         # if quiz in quiz_storage exists already, directly get it
         index, desc = item
-        _description_id = compute_content_hash(desc, prefix="quiz-")
-        if self.quiz_storage.get_by_id(_description_id):
+        _quiz_id = compute_dict_hash({"index": index, "description": desc})
+        if self.quiz_storage.get_by_id(_quiz_id):
             return None
 
         tasks = []
@@ -56,7 +56,7 @@ async def _process_single_quiz(self, item: tuple) -> dict | None:
                 rephrased_text = self.generator.parse_rephrased_text(new_description)
                 quizzes.append((rephrased_text, gt))
             return {
-                "_description_id": _description_id,
+                "_quiz_id": _quiz_id,
                 "description": desc,
                 "index": index,
                 "quizzes": quizzes,
@@ -100,7 +100,7 @@ def quiz(self) -> Iterable[pd.DataFrame]:
                 if new_result:
                     self.quiz_storage.upsert(
                         {
-                            new_result["_description_id"]: {
+                            new_result["_quiz_id"]: {
                                 "description": new_result["description"],
                                 "quizzes": new_result["quizzes"],
                             }