fix(models): add difficulty order

Author: ChenZiHong-Gavin

graphgen/operators/traverse_graph.py

@@ -171,20 +171,21 @@ async def _process_single_batch(
         losses.append(loss)
     q1, q2 = get_loss_tercile(losses)
 
+    difficulty_order = traverse_strategy.difficulty_order
     for i, batch in enumerate(processing_batches):
         if len(batch[1]) == 0:
-            processing_batches[i] = (batch[0], batch[1], "easy")
+            processing_batches[i] = (batch[0], batch[1], difficulty_order[0])
             continue
         loss = sum(edge[2]['loss'] for edge in batch[1]) / len(batch[1])
         if loss < q1:
             # easy
-            processing_batches[i] = (batch[0], batch[1], "easy")
+            processing_batches[i] = (batch[0], batch[1], difficulty_order[0])
         elif loss < q2:
             # medium
-            processing_batches[i] = (batch[0], batch[1], "medium")
+            processing_batches[i] = (batch[0], batch[1], difficulty_order[1])
         else:
             # hard
-            processing_batches[i] = (batch[0], batch[1], "hard")
+            processing_batches[i] = (batch[0], batch[1], difficulty_order[2])
 
     for result in tqdm_async(asyncio.as_completed(
         [_process_single_batch(batch) for batch in processing_batches]

models/strategy/travserse_strategy.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from models.strategy.base_strategy import BaseStrategy
 
@@ -19,6 +19,8 @@ class TraverseStrategy(BaseStrategy):
     edge_sampling: str = "max_loss" # "max_loss" or "min_loss" or "random"
     # 孤立节点的处理策略
     isolated_node_strategy: str = "add" # "add" or "ignore"
+    # 难度顺序 ["easy", "medium", "hard"], ["hard", "medium", "easy"], ["medium", "medium", "medium"]
+    difficulty_order: list = field(default_factory=lambda: ["easy", "medium", "hard"])
 
     def to_yaml(self):
         return {
@@ -29,6 +31,7 @@ def to_yaml(self):
                 "max_tokens": self.max_tokens,
                 "max_depth": self.max_depth,
                 "edge_sampling": self.edge_sampling,
-                "isolated_node_strategy": self.isolated_node_strategy
+                "isolated_node_strategy": self.isolated_node_strategy,
+                "difficulty_order": self.difficulty_order
             }
         }