wip: add protein qa pipeline

Author: ChenZiHong-Gavin

graphgen/bases/datatypes.py

@@ -15,7 +15,7 @@ def from_dict(key: str, data: dict) -> "Chunk":
         return Chunk(
             id=key,
             content=data.get("content", ""),
-            type=data.get("type", "unknown"),
+            type=data.get("type", "text"),
             metadata={k: v for k, v in data.items() if k != "content"},
         )
 

graphgen/configs/protein_qa_config.yaml

@@ -1,5 +1,6 @@
 read:
-  input_file: resources/input_examples/protein_demo.json # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+  input_file: resources/input_examples/protein_qa_demo.txt # input file path, support json, jsonl, txt, pdf. See resources/input_examples for examples
+  anchor_type: protein # get protein information from chunks
 split:
   chunk_size: 1024 # chunk size for text splitting
   chunk_overlap: 100 # chunk overlap for text splitting

graphgen/graphgen.py

@@ -118,16 +118,16 @@ async def insert(self, read_config: Dict, split_config: Dict):
             self.progress_bar,
         )
 
-        _add_chunk_keys = await self.chunks_storage.filter_keys(
-            list(inserting_chunks.keys())
-        )
-        inserting_chunks = {
-            k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
-        }
-
-        if len(inserting_chunks) == 0:
-            logger.warning("All chunks are already in the storage")
-            return
+        # _add_chunk_keys = await self.chunks_storage.filter_keys(
+        #     list(inserting_chunks.keys())
+        # )
+        # inserting_chunks = {
+        #     k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
+        # }
+        #
+        # if len(inserting_chunks) == 0:
+        #     logger.warning("All chunks are already in the storage")
+        #     return
 
         logger.info("[New Chunks] inserting %d chunks", len(inserting_chunks))
         await self.chunks_storage.upsert(inserting_chunks)
@@ -136,6 +136,7 @@ async def insert(self, read_config: Dict, split_config: Dict):
             llm_client=self.synthesizer_llm_client,
             kg_instance=self.graph_storage,
             chunks=[Chunk.from_dict(k, v) for k, v in inserting_chunks.items()],
+            anchor_type=read_config.get("anchor_type", None),
             progress_bar=self.progress_bar,
         )
         if not _add_entities_and_relations:

graphgen/operators/build_kg/build_kg.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import gradio as gr
 
@@ -8,20 +8,23 @@
 from graphgen.utils import logger
 
 from .build_mm_kg import build_mm_kg
+from .build_mo_kg import build_mo_kg
 from .build_text_kg import build_text_kg
 
 
 async def build_kg(
     llm_client: OpenAIClient,
     kg_instance: BaseGraphStorage,
     chunks: List[Chunk],
+    anchor_type: Optional[str] = None,
     progress_bar: gr.Progress = None,
 ):
     """
     Build knowledge graph (KG) and merge into kg_instance
     :param llm_client: Synthesizer LLM model to extract entities and relationships
     :param kg_instance
     :param chunks
+    :param anchor_type: get this type of information from chunks
     :param progress_bar: Gradio progress bar to show the progress of the extraction
     :return:
     """
@@ -49,4 +52,17 @@ async def build_kg(
             chunks=mm_chunks,
             progress_bar=progress_bar,
         )
+
+    if anchor_type is not None:
+        logger.info("Anchoring data based on %s ...", anchor_type)
+        if anchor_type == "protein":
+            await build_mo_kg(
+                llm_client=llm_client,
+                kg_instance=kg_instance,
+                chunks=text_chunks,
+                progress_bar=progress_bar,
+            )
+        else:
+            logger.error("Anchor type %s is not supported yet.", anchor_type)
+
     return kg_instance

graphgen/operators/build_kg/build_mo_kg.py

@@ -1,10 +1,21 @@
+import json
+import re
 from typing import List
 
 import gradio as gr
 
 from graphgen.bases.base_storage import BaseGraphStorage
 from graphgen.bases.datatypes import Chunk
 from graphgen.models import OpenAIClient
+from graphgen.templates import PROTEIN_ANCHOR_PROMPT, PROTEIN_KG_EXTRACTION_PROMPT
+from graphgen.utils import (
+    detect_main_language,
+    handle_single_entity_extraction,
+    handle_single_relationship_extraction,
+    logger,
+    run_concurrent,
+    split_string_by_multi_markers,
+)
 
 
 async def build_mo_kg(
@@ -21,9 +32,89 @@ async def build_mo_kg(
     :param progress_bar: Gradio progress bar to show the progress of the extraction
     :return:
     """
-    # TODO: implement multi-omics KG building logic here
-    print("llm_client:", llm_client)
-    print("kg_instance:", kg_instance)
-    print("chunks:", chunks)
-    print("progress_bar:", progress_bar)
-    return kg_instance
+
+    async def extract_mo_info(chunk: Chunk):
+        content = chunk.content
+        language = detect_main_language(content)
+        prompt = PROTEIN_ANCHOR_PROMPT[language].format(chunk=content)
+        result = await llm_client.generate_answer(prompt)
+        try:
+            json_result = json.loads(result)
+            return json_result
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse JSON from LLM response: %s", result)
+            return {}
+
+    results = await run_concurrent(
+        extract_mo_info,
+        chunks,
+        desc="Extracting multi-omics anchoring information from chunks",
+        unit="chunk",
+        progress_bar=progress_bar,
+    )
+    # Merge results
+    from collections import defaultdict
+
+    bags = defaultdict(set)
+    for item in results:
+        for k, v in item.items():
+            if v is not None and str(v).strip():
+                bags[k].add(str(v).strip())
+
+    merged = {
+        k: " | ".join(sorted(v)) if len(v) > 1 else next(iter(v))
+        for k, v in bags.items()
+    }
+
+    # TODO: search database for more info
+    # try:
+    #     search_results = await search(merged["Protein accession or ID"])
+    # except Exception as e:
+    #     logger.warning("Failed to search for protein info: %s", e)
+    #     search_results = {}
+
+    # 组织成文本
+    mo_text = "\n".join([f"{k}: {v}" for k, v in merged.items()])
+    lang = detect_main_language(mo_text)
+    prompt = PROTEIN_KG_EXTRACTION_PROMPT[lang].format(
+        input_text=mo_text,
+        **PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"],
+    )
+    kg_output = await llm_client.generate_answer(prompt)
+
+    logger.debug("Image chunk extraction result: %s", kg_output)
+
+    # parse the result
+    records = split_string_by_multi_markers(
+        kg_output,
+        [
+            PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["record_delimiter"],
+            PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["completion_delimiter"],
+        ],
+    )
+
+    print(records)
+    raise NotImplementedError
+
+    nodes = defaultdict(list)
+    edges = defaultdict(list)
+
+    for record in records:
+        match = re.search(r"\((.*)\)", record)
+        if not match:
+            continue
+        inner = match.group(1)
+
+        attributes = split_string_by_multi_markers(
+            inner, [PROTEIN_KG_EXTRACTION_PROMPT["FORMAT"]["tuple_delimiter"]]
+        )
+
+        entity = await handle_single_entity_extraction(attributes, "temp")
+        if entity is not None:
+            nodes[entity["entity_name"]].append(entity)
+            continue
+
+        relation = await handle_single_relationship_extraction(attributes, "temp")
+        if relation is not None:
+            key = (relation["src_id"], relation["tgt_id"])
+            edges[key].append(relation)