Refactor logging setup and clean up whitespace in various files. Adjusted formatting in base_operator.py, base_searcher.py, and anchor_bfs_partitioner.py for consistency. Removed unnecessary blank lines in json_reader.py, rnacentral_searcher.py, and other template files to improve readability.

Author: CHERRY-ui8

graphgen/bases/base_operator.py

@@ -30,7 +30,7 @@ def __init__(self, working_dir: str = "cache", op_name: str = None):
         log_file = os.path.join(log_dir, f"{self.op_name}_{worker_id_short}.log")
 
         self.logger = set_logger(
-            log_file=log_file, name=f"{self.op_name}.{worker_id_short}", 
+            log_file=log_file, name=f"{self.op_name}.{worker_id_short}",
             console_level=logging.ERROR, force=True
         )
 

graphgen/bases/base_searcher.py

@@ -24,7 +24,7 @@ def __init__(self, working_dir: str = "cache"):
         log_file = os.path.join(log_dir, f"{searcher_name}.log")
 
         self.logger = set_logger(
-            log_file=log_file, name=searcher_name, 
+            log_file=log_file, name=searcher_name,
             console_level=logging.ERROR, force=True
         )
 

graphgen/models/partitioner/anchor_bfs_partitioner.py

@@ -18,7 +18,7 @@ class AnchorBFSPartitioner(BFSPartitioner):
     2. Expand the community using BFS until the max unit size is reached.(A unit is a node or an edge.)
     3. Non-anchor units can only be "pulled" into a community and never become seeds themselves.
     For example, for VQA tasks, we may want to use image nodes as anchors and expand to nearby text nodes and edges.
-    
+
     Supports multiple anchor types for multi-omics data: anchor_type can be a single string or a list of strings.
     When a list is provided, nodes matching any of the types in the list can serve as anchors.
     """
@@ -79,39 +79,39 @@ def _pick_anchor_ids(
 
         anchor_ids: Set[str] = set()
         anchor_types_lower = [at.lower() for at in self.anchor_types]
-        
+
         for node_id, meta in nodes:
             # Check if node matches any of the anchor types
             matched = False
-            
+
             # Check 1: entity_type (for image, etc.)
             node_type = str(meta.get("entity_type", "")).lower()
             for anchor_type_lower in anchor_types_lower:
                 if anchor_type_lower in node_type:
                     anchor_ids.add(node_id)
                     matched = True
                     break
-            
+
             if matched:
                 continue
-            
+
             # Check 2: molecule_type (for omics data: dna, rna, protein)
             molecule_type = str(meta.get("molecule_type", "")).lower()
             if molecule_type in anchor_types_lower:
                 anchor_ids.add(node_id)
                 continue
-            
+
             # Check 3: source_id prefix (for omics data: dna-, rna-, protein-)
             source_id = str(meta.get("source_id", "")).lower()
             for anchor_type_lower in anchor_types_lower:
                 if source_id.startswith(f"{anchor_type_lower}-"):
                     anchor_ids.add(node_id)
                     matched = True
                     break
-            
+
             if matched:
                 continue
-            
+
             # Check 4: Check if source_id contains multiple IDs separated by <SEP>
             if "<sep>" in source_id:
                 source_ids = source_id.split("<sep>")
@@ -124,7 +124,7 @@ def _pick_anchor_ids(
                             break
                     if matched:
                         break
-        
+
         return anchor_ids
 
     @staticmethod

graphgen/models/reader/json_reader.py

@@ -48,13 +48,13 @@ def read_stream(self, file_path: str) -> Iterator[Dict[str, Any]]:
         """
         Stream read JSONL files line by line without loading entire file into memory.
         Returns an iterator that yields filtered documents.
-        
+
         :param file_path: Path to the JSONL file.
         :return: Iterator of dictionaries containing the data.
         """
         if not file_path.endswith(".jsonl"):
             raise ValueError("read_stream only supports JSONL files, not JSON files")
-        
+
         with open(file_path, "r", encoding="utf-8") as f:
             for line in f:
                 try:
@@ -64,7 +64,7 @@ def read_stream(self, file_path: str) -> Iterator[Dict[str, Any]]:
                         raise ValueError(
                             f"Missing '{self.text_column}' in document: {doc}"
                         )
-                    
+
                     # Apply filtering logic inline (similar to BaseReader.filter)
                     if doc.get("type") == "text":
                         content = doc.get(self.text_column, "").strip()

graphgen/models/searcher/db/rnacentral_searcher.py

@@ -35,9 +35,9 @@ class RNACentralSearch(BaseSearcher):
     """
 
     def __init__(
-        self, 
-        use_local_blast: bool = False, 
-        local_blast_db: str = "rna_db", 
+        self,
+        use_local_blast: bool = False,
+        local_blast_db: str = "rna_db",
         api_timeout: int = 30,
         blast_num_threads: int = 4,
         working_dir: str = "cache",
@@ -49,7 +49,7 @@ def __init__(
         self.local_blast_db = local_blast_db
         self.api_timeout = api_timeout
         self.blast_num_threads = blast_num_threads  # Number of threads for BLAST search
-        
+
         if self.use_local_blast and not os.path.isfile(f"{self.local_blast_db}.nhr"):
             self.logger.error("Local BLAST database files not found. Please check the path.")
             self.use_local_blast = False
@@ -254,22 +254,22 @@ def _local_blast(self, seq: str, threshold: float) -> Optional[str]:
                 "-num_threads", str(self.blast_num_threads),
                 "-outfmt", "6 sacc"  # Only accession, tab-separated
             ]
-            self.logger.debug("Running local blastn for RNA (threads=%d): %s", 
+            self.logger.debug("Running local blastn for RNA (threads=%d): %s",
                         self.blast_num_threads, " ".join(cmd))
-            
+
             # Run BLAST with timeout to avoid hanging
             try:
                 out = subprocess.check_output(
-                    cmd, 
-                    text=True, 
+                    cmd,
+                    text=True,
                     timeout=300,  # 5 minute timeout for BLAST search
                     stderr=subprocess.DEVNULL  # Suppress BLAST warnings to reduce I/O
                 ).strip()
             except subprocess.TimeoutExpired:
                 self.logger.warning("BLAST search timed out after 5 minutes for sequence")
                 os.remove(tmp_name)
                 return None
-            
+
             os.remove(tmp_name)
             return out.split("\n", maxsplit=1)[0] if out else None
         except Exception as exc:
@@ -378,7 +378,7 @@ async def search(self, query: str, threshold: float = 0.1, **kwargs) -> Optional
         # check if RNA sequence (AUCG or ATCG characters, contains U or T)
         # Note: Sequences with T are also RNA sequences
         is_rna_sequence = query.startswith(">") or (
-            re.fullmatch(r"[AUCGTN\s]+", query, re.I) and 
+            re.fullmatch(r"[AUCGTN\s]+", query, re.I) and
             ("U" in query.upper() or "T" in query.upper())
         )
         if is_rna_sequence:

graphgen/models/splitter/sequence_splitter.py

@@ -20,7 +20,7 @@ def __init__(
     ):
         """
         Initialize sequence splitter.
-        
+
         :param chunk_size: Maximum length of each chunk (in sequence characters)
         :param chunk_overlap: Number of characters to overlap between chunks
         :param length_function: Function to calculate length (default: len)
@@ -37,83 +37,83 @@ def __init__(
     def split_text(self, text: str) -> List[str]:
         """
         Split a sequence into chunks of fixed size with overlap.
-        
+
         :param text: The sequence to split (may include FASTA header)
         :return: List of sequence chunks
         """
         # Remove FASTA header if present
         sequence = self._extract_sequence(text)
-        
+
         if not sequence:
             logger.warning("Empty sequence provided to SequenceSplitter")
             return []
-        
+
         # If sequence is shorter than chunk_size, return as single chunk
         if len(sequence) <= self.chunk_size:
             return [sequence]
-        
+
         chunks = []
         start = 0
         step = self.chunk_size - self.chunk_overlap
-        
+
         while start < len(sequence):
             end = min(start + self.chunk_size, len(sequence))
             chunk = sequence[start:end]
             chunks.append(chunk)
             start += step
-            
+
             # Avoid infinite loop if step is 0 or negative
             if step <= 0:
                 break
-        
+
         return chunks
 
     @staticmethod
     def _extract_sequence(text: str) -> str:
         """
         Extract sequence from text, removing FASTA headers and whitespace.
-        
+
         :param text: Input text (may contain FASTA header)
         :return: Clean sequence string
         """
         # Remove FASTA header lines (lines starting with >)
         lines = text.split("\n")
         sequence_lines = [line for line in lines if not line.strip().startswith(">")]
-        
+
         # Join and remove whitespace
         sequence = "".join(sequence_lines)
         sequence = re.sub(r"\s+", "", sequence)
-        
+
         return sequence.upper()  # Normalize to uppercase
 
     @staticmethod
     def detect_sequence_type(sequence: str) -> Optional[str]:
         """
         Detect the type of sequence (DNA, RNA, or protein).
-        
+
         :param sequence: The sequence string
         :return: "dna", "rna", "protein", or None if cannot determine
         """
         # Remove FASTA header and whitespace
         clean_seq = SequenceSplitter._extract_sequence(sequence)
-        
+
         if not clean_seq:
             return None
-        
+
         # Check for protein-specific amino acids
         protein_chars = set("EFILPQXZ")  # Amino acids not in DNA/RNA
         if any(char in clean_seq for char in protein_chars):
             return "protein"
-        
+
         # Check for RNA-specific character (U)
         if "U" in clean_seq.upper():
             return "rna"
-        
+
         # Check if contains only DNA/RNA characters (A, T, G, C, N)
         dna_rna_chars = set("ATGCUN")
         if all(char.upper() in dna_rna_chars for char in clean_seq):
             # Default to DNA if ambiguous (could be DNA or RNA without U)
             return "dna"
-        
+
         # If contains other characters, might be protein
         return "protein"

graphgen/models/storage/kv/json_storage.py

@@ -54,7 +54,7 @@ def iter_items(self) -> Iterator[Tuple[str, dict]]:
     def get_batch(self, keys: list[str]) -> dict[str, dict]:
         """
         Get a batch of items by their keys.
-        
+
         :param keys: List of keys to retrieve.
         :return: Dictionary of {key: value} for the requested keys.
         """
@@ -63,7 +63,7 @@ def get_batch(self, keys: list[str]) -> dict[str, dict]:
     def iter_batches(self, batch_size: int = 10000) -> Iterator[dict[str, dict]]:
         """
         Iterate over items in batches to avoid loading everything into memory.
-        
+
         :param batch_size: Number of items per batch.
         :return: Iterator of dictionaries, each containing up to batch_size items.
         """

graphgen/operators/build_kg/build_omics_kg.py

@@ -15,7 +15,7 @@ def build_omics_kg(
 ):
     """
     Build knowledge graph from multi-omics chunks (DNA, RNA, protein).
-    
+
     :param llm_client: Synthesizer LLM model to extract entities and relationships
     :param kg_instance: Graph storage instance
     :param chunks: List of omics chunks

graphgen/templates/description_rephrasing.py

@@ -19,7 +19,7 @@
 ################
 -Real Data-
 ################
-Input: 
+Input:
 {input_sentence}
 ################
 Please directly output the rewritten sentence without any additional information.

graphgen/templates/extraction/schema_guided_extraction.py

@@ -7,7 +7,7 @@
 3. Present your findings in JSON format as specified below.
 
 Important Notes:
-- Extract only relevant information. 
+- Extract only relevant information.
 - Consider the context of the entire document when determining relevance.
 - Do not be verbose, only respond with the correct format and information.
 - Some docs may have multiple relevant excerpts -- include all that apply.