Remove @DataClass from all subclasses following generator pattern

Co-authored-by: ChenZiHong-Gavin <[email protected]>

Author: Copilot

graphgen/models/evaluator/length_evaluator.py

@@ -1,16 +1,13 @@
-from dataclasses import dataclass
-
 from graphgen.bases.datatypes import QAPair
 from graphgen.models.evaluator.base_evaluator import BaseEvaluator
 from graphgen.models.tokenizer import Tokenizer
 from graphgen.utils import create_event_loop
 
 
-@dataclass
 class LengthEvaluator(BaseEvaluator):
-    tokenizer_name: str = "cl100k_base"
-
-    def __post_init__(self):
+    def __init__(self, tokenizer_name: str = "cl100k_base", max_concurrent: int = 100):
+        super().__init__(max_concurrent)
+        self.tokenizer_name = tokenizer_name
         self.tokenizer = Tokenizer(model_name=self.tokenizer_name)
 
     async def evaluate_single(self, pair: QAPair) -> float:

graphgen/models/evaluator/mtld_evaluator.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass, field
 from typing import Set
 
 from graphgen.bases.datatypes import QAPair
@@ -8,18 +7,15 @@
 nltk_helper = NLTKHelper()
 
 
-@dataclass
 class MTLDEvaluator(BaseEvaluator):
     """
     衡量文本词汇多样性的指标
     """
 
-    stopwords_en: Set[str] = field(
-        default_factory=lambda: set(nltk_helper.get_stopwords("english"))
-    )
-    stopwords_zh: Set[str] = field(
-        default_factory=lambda: set(nltk_helper.get_stopwords("chinese"))
-    )
+    def __init__(self, max_concurrent: int = 100):
+        super().__init__(max_concurrent)
+        self.stopwords_en: Set[str] = set(nltk_helper.get_stopwords("english"))
+        self.stopwords_zh: Set[str] = set(nltk_helper.get_stopwords("chinese"))
 
     async def evaluate_single(self, pair: QAPair) -> float:
         loop = create_event_loop()

graphgen/models/kg_builder/light_rag_kg_builder.py

@@ -1,6 +1,5 @@
 import re
 from collections import Counter, defaultdict
-from dataclasses import dataclass
 from typing import Dict, List, Tuple
 
 from graphgen.bases import BaseGraphStorage, BaseKGBuilder, BaseLLMClient, Chunk
@@ -15,10 +14,10 @@
 )
 
 
-@dataclass
 class LightRAGKGBuilder(BaseKGBuilder):
-    llm_client: BaseLLMClient = None
-    max_loop: int = 3
+    def __init__(self, llm_client: BaseLLMClient, max_loop: int = 3):
+        super().__init__(llm_client)
+        self.max_loop = max_loop
 
     async def extract(
         self, chunk: Chunk

graphgen/models/kg_builder/mm_kg_builder.py

@@ -1,6 +1,5 @@
 import re
 from collections import defaultdict
-from dataclasses import dataclass
 from typing import Dict, List, Tuple
 
 from graphgen.bases import BaseLLMClient, Chunk
@@ -16,11 +15,7 @@
 from .light_rag_kg_builder import LightRAGKGBuilder
 
 
-@dataclass
 class MMKGBuilder(LightRAGKGBuilder):
-    llm_client: BaseLLMClient = None
-    max_loop: int = 3
-
     async def extract(
         self, chunk: Chunk
     ) -> Tuple[Dict[str, List[dict]], Dict[Tuple[str, str], List[dict]]]:

graphgen/models/partitioner/bfs_partitioner.py

@@ -1,6 +1,5 @@
 import random
 from collections import deque
-from dataclasses import dataclass
 from typing import Any, List
 
 from graphgen.bases import BaseGraphStorage, BasePartitioner
@@ -10,7 +9,6 @@
 EDGE_UNIT: str = "e"
 
 
-@dataclass
 class BFSPartitioner(BasePartitioner):
     """
     BFS partitioner that partitions the graph into communities of a fixed size.

graphgen/models/partitioner/dfs_partitioner.py

@@ -1,5 +1,4 @@
 import random
-from dataclasses import dataclass
 from typing import Any, List
 
 from graphgen.bases import BaseGraphStorage, BasePartitioner
@@ -9,7 +8,6 @@
 EDGE_UNIT: str = "e"
 
 
-@dataclass
 class DFSPartitioner(BasePartitioner):
     """
     DFS partitioner that partitions the graph into communities of a fixed size.

graphgen/models/partitioner/ece_partitioner.py

@@ -1,6 +1,5 @@
 import asyncio
 import random
-from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Set, Tuple
 
 from tqdm.asyncio import tqdm as tqdm_async
@@ -13,7 +12,6 @@
 EDGE_UNIT: str = "e"
 
 
-@dataclass
 class ECEPartitioner(BFSPartitioner):
     """
     ECE partitioner that partitions the graph into communities based on Expected Calibration Error (ECE).

graphgen/models/partitioner/leiden_partitioner.py

@@ -1,5 +1,4 @@
 from collections import defaultdict
-from dataclasses import dataclass
 from typing import Any, Dict, List, Set, Tuple
 
 import igraph as ig
@@ -9,7 +8,6 @@
 from graphgen.bases.datatypes import Community
 
 
-@dataclass
 class LeidenPartitioner(BasePartitioner):
     """
     Leiden partitioner that partitions the graph into communities using the Leiden algorithm.

graphgen/models/tokenizer/__init__.py

@@ -1,4 +1,3 @@
-from dataclasses import dataclass, field
 from typing import List
 
 from graphgen.bases import BaseTokenizer
@@ -30,16 +29,13 @@ def get_tokenizer_impl(tokenizer_name: str = "cl100k_base") -> BaseTokenizer:
     )
 
 
-@dataclass
 class Tokenizer(BaseTokenizer):
     """
     Encapsulates different tokenization implementations based on the specified model name.
     """
 
-    model_name: str = "cl100k_base"
-    _impl: BaseTokenizer = field(init=False, repr=False)
-
-    def __post_init__(self):
+    def __init__(self, model_name: str = "cl100k_base"):
+        super().__init__(model_name)
         if not self.model_name:
             raise ValueError("TOKENIZER_MODEL must be specified in the ENV variables.")
         self._impl = get_tokenizer_impl(self.model_name)

graphgen/models/tokenizer/hf_tokenizer.py

@@ -1,16 +1,13 @@
-from dataclasses import dataclass
 from typing import List
 
 from transformers import AutoTokenizer
 
 from graphgen.bases import BaseTokenizer
 
 
-@dataclass
 class HFTokenizer(BaseTokenizer):
-    model_name: str = "cl100k_base"
-
-    def __post_init__(self):
+    def __init__(self, model_name: str = "cl100k_base"):
+        super().__init__(model_name)
         self.enc = AutoTokenizer.from_pretrained(self.model_name)
 
     def encode(self, text: str) -> List[int]: