feat(graphgen): calculate node loss

Author: ChenZiHong-Gavin

generate.py

@@ -73,7 +73,7 @@
 
     graph_gen.insert(data, args.data_type)
 
-    graph_gen.quiz(max_samples=3)
+    graph_gen.quiz(max_samples=2)
 
     graph_gen.judge(re_judge=False)
 

graphgen/graphgen.py

@@ -10,7 +10,7 @@
 from models import Chunk, JsonKVStorage, OpenAIModel, NetworkXStorage, WikiSearch, Tokenizer, TraverseStrategy
 from models.storage.base_storage import StorageNameSpace
 from utils import create_event_loop, logger, compute_content_hash
-from .operators import extract_kg, search_wikipedia, quiz_relations, judge_relations, traverse_graph_by_edge
+from .operators import extract_kg, search_wikipedia, quiz, judge_statement, traverse_graph_by_edge
 
 
 sys_path = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -169,15 +169,15 @@ def quiz(self, max_samples=1):
         loop.run_until_complete(self.async_quiz(max_samples))
 
     async def async_quiz(self, max_samples=1):
-        await quiz_relations(self.teacher_llm_client, self.graph_storage, self.rephrase_storage, max_samples)
+        await quiz(self.teacher_llm_client, self.graph_storage, self.rephrase_storage, max_samples)
         await self.rephrase_storage.index_done_callback()
 
     def judge(self, re_judge=False):
         loop = create_event_loop()
         loop.run_until_complete(self.async_judge(re_judge))
 
     async def async_judge(self, re_judge=False):
-        _update_relations = await judge_relations(self.student_llm_client, self.graph_storage,
+        _update_relations = await judge_statement(self.student_llm_client, self.graph_storage,
                                                   self.rephrase_storage, re_judge)
         await _update_relations.index_done_callback()
 

graphgen/operators/__init__.py

@@ -1,13 +1,13 @@
 from .extract_kg import extract_kg
-from .quiz_relations import quiz_relations
-from .judge_relations import judge_relations
+from .quiz import quiz
+from .judge import judge_statement
 from .search_wikipedia import search_wikipedia
 from .traverse_graph import traverse_graph_by_edge
 
 __all__ = [
     "extract_kg",
-    "quiz_relations",
-    "judge_relations",
+    "quiz",
+    "judge_statement",
     "search_wikipedia",
     "traverse_graph_by_edge"
 ]

graphgen/operators/judge_relations.py graphgen/operators/judge.pygraphgen/operators/judge_relations.py renamed to graphgen/operators/judge.py

@@ -6,14 +6,14 @@
 from templates import STATEMENT_JUDGEMENT_PROMPT
 
 
-async def judge_relations(
+async def judge_statement(
         student_llm_client: OpenAIModel,
         graph_storage: NetworkXStorage,
         rephrase_storage: JsonKVStorage,
         re_judge: bool = False,
         max_concurrent: int = 1000) -> NetworkXStorage:
     """
-    Get all edges and judge them
+    Get all edges and nodes and judge them
 
     :param student_llm_client: judge the statements to get comprehension loss
     :param graph_storage: graph storage instance
@@ -74,4 +74,52 @@ async def _judge_single_relation(
     ):
         results.append(await result)
 
+    async def _judge_single_entity(
+        node: tuple,
+    ):
+        async with semaphore:
+            node_id = node[0]
+            node_data = node[1]
+
+            if (not re_judge) and "loss" in node_data and node_data["loss"] is not None:
+                logger.info("Node %s already judged, loss: %s, skip", node_id, node_data["loss"])
+                return node_id, node_data
+
+            description = node_data["description"]
+
+            try:
+                descriptions = await rephrase_storage.get_by_id(description)
+                assert descriptions is not None
+
+                judgements = []
+                gts = [gt for _, gt in descriptions]
+                for description, gt in descriptions:
+                    judgement = await student_llm_client.generate_topk_per_token(
+                        STATEMENT_JUDGEMENT_PROMPT['TEMPLATE'].format(statement=description)
+                    )
+                    judgements.append(judgement[0].top_candidates)
+
+                loss = yes_no_loss_entropy(judgements, gts)
+
+                logger.info("Node %s description: %s loss: %s", node_id, description, loss)
+
+                node_data["loss"] = loss
+            except Exception as e: # pylint: disable=broad-except
+                logger.error("Error in judging entity %s: %s", node_id, e)
+                logger.info("Use default loss 0.1")
+                node_data["loss"] = -math.log(0.1)
+
+            await graph_storage.update_node(node_id, node_data)
+            return node_id, node_data
+
+    nodes = await graph_storage.get_all_nodes()
+
+    results = []
+    for result in tqdm_async(
+            asyncio.as_completed([_judge_single_entity(node) for node in nodes]),
+            total=len(nodes),
+            desc="Judging entities"
+    ):
+        results.append(await result)
+
     return graph_storage

graphgen/operators/quiz_relations.py graphgen/operators/quiz.pygraphgen/operators/quiz_relations.py renamed to graphgen/operators/quiz.py

@@ -7,7 +7,7 @@
 from templates import DESCRIPTION_REPHRASING_PROMPT
 
 
-async def quiz_relations(
+async def quiz(
         teacher_llm_client: OpenAIModel,
         graph_storage: NetworkXStorage,
         rephrase_storage: JsonKVStorage,
@@ -26,16 +26,12 @@ async def quiz_relations(
 
     semaphore = asyncio.Semaphore(max_concurrent)
 
-    async def _quiz_single_relation(
-        edge: tuple,
+    async def _process_single_quiz(
         des: str,
         prompt: str,
         gt: str
     ):
         async with semaphore:
-            source_id = edge[0]
-            target_id = edge[1]
-
             try:
                 # 如果在rephrase_storage中已经存在，直接取出
                 descriptions = await rephrase_storage.get_by_id(des)
@@ -49,11 +45,12 @@ async def _quiz_single_relation(
                 return  {des: [(new_description, gt)]}
 
             except Exception as e: # pylint: disable=broad-except
-                logger.error("Error when quizzing edge %s -> %s: %s", source_id, target_id, e)
+                logger.error("Error when quizzing description %s: %s", des, e)
                 return None
 
 
     edges = await graph_storage.get_all_edges()
+    nodes = await graph_storage.get_all_nodes()
 
     results = defaultdict(list)
     tasks = []
@@ -68,19 +65,36 @@ async def _quiz_single_relation(
         for i in range(max_samples):
             if i > 0:
                 tasks.append(
-                    _quiz_single_relation(edge, description,
+                    _process_single_quiz(description,
                                           DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
                                               input_sentence=description), 'yes')
                 )
-            tasks.append(_quiz_single_relation(edge, description,
+            tasks.append(_process_single_quiz(description,
                                               DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
                                                   input_sentence=description), 'no'))
 
+    for node in nodes:
+        node_data = node[1]
+        description = node_data["description"]
+        language = "English" if detect_main_language(description) == "en" else "Chinese"
+
+        results[description] = [(description, 'yes')]
+
+        for i in range(max_samples):
+            if i > 0:
+                tasks.append(
+                    _process_single_quiz(description,
+                                          DESCRIPTION_REPHRASING_PROMPT[language]['TEMPLATE'].format(
+                                              input_sentence=description), 'yes')
+                )
+            tasks.append(_process_single_quiz(description,
+                                              DESCRIPTION_REPHRASING_PROMPT[language]['ANTI_TEMPLATE'].format(
+                                                  input_sentence=description), 'no'))
 
     for result in tqdm_async(
             asyncio.as_completed(tasks),
             total=len(tasks),
-            desc="Quizzing relations"
+            desc="Quizzing descriptions"
     ):
         new_result = await result
         if new_result:
@@ -91,4 +105,5 @@ async def _quiz_single_relation(
         results[key] = list(set(value))
         await rephrase_storage.upsert({key: results[key]})
 
+
     return rephrase_storage

graphgen/operators/split_graph.py

@@ -260,6 +260,7 @@ async def get_cached_node_info(node_id: str) -> dict:
 
         processing_batches.append((_process_nodes, _process_edges))
 
+    l
     # isolate nodes
     isolated_node_strategy = traverse_strategy.isolated_node_strategy
     if isolated_node_strategy == "add":

graphgen/operators/traverse_graph.py

@@ -57,6 +57,10 @@ def get_loss_tercile(losses: list) -> (float, float):
 
     return losses[q1_index], losses[q2_index]
 
+def get_average_loss(batch: tuple) -> float:
+    return sum(edge[2]['loss'] for edge in batch[1]) + sum(node['loss'] for node in batch[0]) / \
+           (len(batch[0]) + len(batch[1]))
+
 async def traverse_graph_by_edge(
     llm_client: OpenAIModel,
     tokenizer: Tokenizer,
@@ -114,8 +118,6 @@ async def _process_single_batch(
         _process_batch: tuple
     ) -> dict:
         async with semaphore:
-            losses = [(edge[0], edge[1], edge[2]['loss']) for edge in _process_batch[1]]
-
             context = await _process_nodes_and_edges(
                 _process_batch[0],
                 _process_batch[1],
@@ -145,14 +147,14 @@ async def _process_single_batch(
                 compute_content_hash(context): {
                     "question": question,
                     "answer": context,
-                    "losses": losses,
+                    "loss": get_average_loss(_process_batch),
                     "difficulty": _process_batch[2],
                 }
             }
 
     results = {}
     edges = list(await graph_storage.get_all_edges())
-    nodes = await graph_storage.get_all_nodes()
+    nodes = list(await graph_storage.get_all_nodes())
 
     edges, nodes = await _pre_tokenize(graph_storage, tokenizer, edges, nodes)
 
@@ -165,18 +167,13 @@ async def _process_single_batch(
 
     losses = []
     for batch in processing_batches:
-        if len(batch[1]) == 0:
-            continue
-        loss = sum(edge[2]['loss'] for edge in batch[1]) / len(batch[1])
+        loss = get_average_loss(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], difficulty_order[0])
-            continue
-        loss = sum(edge[2]['loss'] for edge in batch[1]) / len(batch[1])
+        loss = get_average_loss(batch)
         if loss < q1:
             # easy
             processing_batches[i] = (batch[0], batch[1], difficulty_order[0])

judge.py

@@ -0,0 +1,40 @@
+import os
+import argparse
+import asyncio
+from dotenv import load_dotenv
+
+from models import NetworkXStorage, JsonKVStorage, OpenAIModel
+from graphgen.operators import judge_relations
+
+sys_path = os.path.abspath(os.path.dirname(__file__))
+
+load_dotenv()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--output', type=str, default='cache/output/new_graph.graphml', help='path to save output')
+
+    args = parser.parse_args()
+
+    llm_client = OpenAIModel(
+        model_name=os.getenv("STUDENT_MODEL"),
+        api_key=os.getenv("STUDENT_API_KEY"),
+        base_url=os.getenv("STUDENT_BASE_URL")
+    )
+
+    graph_storage = NetworkXStorage(
+        os.path.join(sys_path, "cache"),
+        namespace="graph"
+    )
+
+    rephrase_storage = JsonKVStorage(
+        os.path.join(sys_path, "cache"),
+        namespace="rephrase"
+    )
+
+    new_graph = asyncio.run(judge_relations(llm_client, graph_storage, rephrase_storage, re_judge=True))
+
+    graph_file = asyncio.run(graph_storage.get_graph())
+
+    new_graph.write_nx_graph(graph_file, args.output)

models/storage/base_storage.py

@@ -1,8 +1,6 @@
-import numpy as np
-
 from dataclasses import dataclass
-from models.embed.embedding import EmbeddingFunc
 from typing import Union, Generic, TypeVar
+from models.embed.embedding import EmbeddingFunc
 
 T = TypeVar("T")
 
@@ -95,6 +93,3 @@ async def upsert_edge(
 
     async def delete_node(self, node_id: str):
         raise NotImplementedError
-
-    async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]:
-        raise NotImplementedError("Node embedding is not used in lightrag.")

models/strategy/travserse_strategy.py

@@ -8,19 +8,19 @@ class TraverseStrategy(BaseStrategy):
     # 最大边数和最大token数方法中选择一个生效
     expand_method: str = "max_tokens" # "max_width" or "max_tokens"
     # 单向拓展还是双向拓展
-    bidirectional: bool = False
+    bidirectional: bool = True
     # 每个方向拓展的最大边数
     max_extra_edges: int = 5
     # 最长token数
-    max_tokens: int = 1024
+    max_tokens: int = 512
     # 每个方向拓展的最大深度
-    max_depth: int = 3
+    max_depth: int = 5
     # 同一层中选边的策略（如果是双向拓展，同一层指的是两边连接的边的集合）
     edge_sampling: str = "max_loss" # "max_loss" or "min_loss" or "random"
     # 孤立节点的处理策略
-    isolated_node_strategy: str = "ignore" # "add" or "ignore"
+    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"])
+    difficulty_order: list = field(default_factory=lambda: ["medium", "medium", "medium"])
 
     def to_yaml(self):
         return {