debase 0.4.1__tar.gz → 0.4.3__tar.gz

debase-0.4.3/.gitignore

@@ -0,0 +1,177 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+Pipfile.lock
+
+# poetry
+poetry.lock
+
+# pdm
+.pdm.toml
+
+# PEP 582
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+.idea/
+
+# VS Code
+.vscode/
+
+# macOS
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Windows
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+*.stackdump
+[Dd]esktop.ini
+$RECYCLE.BIN/
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+*.lnk
+
+# Linux
+*~
+
+# Temporary files
+*.tmp
+*.temp
+*.log
+.temp_*/
+.cache/
+
+# DEBase specific
+enzyme_pipeline*.log
+temp_merged_input.csv
+*.egg-info/
+
+# Project data and examples
+data/
+examples/
+!examples/test.csv  # Keep test.csv as example output
+
+# Cache files
+*.pkl
+*_cache.pkl
+
+# Large database files
+*.db
+
+# PDFs and Excel files
+*.pdf
+*.xlsx
+
+# Backup files
+*_backup.py
+lineage_format_backup.py
+
+# Temporary directories
+.temp_*
+enzyme_analysis_*

{debase-0.4.1/src/debase.egg-info → debase-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: debase
-Version: 0.4.1
+Version: 0.4.3
 Summary: Enzyme lineage analysis and sequence extraction package
 Home-page: https://github.com/YuemingLong/DEBase
 Author: DEBase Team

debase-0.4.3/environment.yml

@@ -0,0 +1,21 @@
+name: debase
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.9
+  - pandas>=1.0.0
+  - numpy>=1.19.0
+  - matplotlib>=3.3.0
+  - seaborn>=0.11.0
+  - jupyter>=1.0.0
+  - jupyterlab>=3.0.0
+  - openpyxl>=3.0.0
+  - biopython>=1.78
+  - requests>=2.25.0
+  - tqdm>=4.60.0
+  - rdkit>=2020.03.1
+  - pip
+  - pip:
+    - PyMuPDF>=1.18.0
+    - google-generativeai>=0.3.0

{debase-0.4.1 → debase-0.4.3}/src/debase/_version.py

@@ -1,3 +1,3 @@
 """Version information."""
 
-__version__ = "0.4.1"
+__version__ = "0.4.3"

{debase-0.4.1 → debase-0.4.3}/src/debase/cleanup_sequence.py

@@ -11,6 +11,7 @@ Usage:
 
 import argparse
 import logging
+import os
 import re
 import sys
 from dataclasses import dataclass, field
@@ -19,11 +20,20 @@ from typing import Dict, List, Optional, Set, Tuple, Union
 
 import pandas as pd
 
+try:
+    import google.generativeai as genai  # type: ignore
+    GEMINI_OK = True
+except ImportError:  # pragma: no cover
+    GEMINI_OK = False
+
 
 # === 1. CONFIGURATION & CONSTANTS === ----------------------------------------
 
 VALID_AMINO_ACIDS = set("ACDEFGHIKLMNPQRSTVWY*")  # Include * for stop codons
 
+# Gemini API configuration
+GEMINI_API_KEY: str = os.environ.get("GEMINI_API_KEY", "")
+
 # Configure module logger
 log = logging.getLogger(__name__)
 
@@ -565,7 +575,136 @@ class SequenceGenerator:
         return None
 
 
-# === 7. MAIN PROCESSOR === ---------------------------------------------------
+# === 7. GEMINI PARENT IDENTIFICATION === ------------------------------------
+
+def identify_parents_with_gemini(df: pd.DataFrame) -> pd.DataFrame:
+    """Use Gemini API to identify parent enzymes for entries with missing parent information."""
+    if not GEMINI_OK:
+        log.warning("Gemini API not available (missing google.generativeai). Skipping parent identification.")
+        return df
+    
+    if not GEMINI_API_KEY:
+        log.warning("GEMINI_API_KEY not set. Skipping parent identification.")
+        return df
+    
+    try:
+        genai.configure(api_key=GEMINI_API_KEY)
+        model = genai.GenerativeModel('gemini-1.5-flash')
+    except Exception as e:
+        log.warning(f"Failed to configure Gemini API: {e}. Skipping parent identification.")
+        return df
+    
+    # Find entries with empty sequences but missing parent information
+    entries_needing_parents = []
+    for idx, row in df.iterrows():
+        protein_seq = str(row.get("protein_sequence", "")).strip()
+        parent_id = str(row.get("parent_enzyme_id", "")).strip()
+        
+        # Only process entries that have empty sequences AND no parent info
+        if (not protein_seq or protein_seq == "nan") and (not parent_id or parent_id == "nan"):
+            enzyme_id = str(row.get("enzyme_id", ""))
+            campaign_id = str(row.get("campaign_id", ""))
+            generation = str(row.get("generation", ""))
+            
+            entries_needing_parents.append({
+                "idx": idx,
+                "enzyme_id": enzyme_id,
+                "campaign_id": campaign_id,
+                "generation": generation
+            })
+    
+    if not entries_needing_parents:
+        log.info("No entries need parent identification from Gemini")
+        return df
+    
+    log.info(f"Found {len(entries_needing_parents)} entries needing parent identification. Querying Gemini...")
+    
+    # Create a lookup of all available enzyme IDs for context
+    available_enzymes = {}
+    for idx, row in df.iterrows():
+        enzyme_id = str(row.get("enzyme_id", ""))
+        campaign_id = str(row.get("campaign_id", ""))
+        protein_seq = str(row.get("protein_sequence", "")).strip()
+        generation = str(row.get("generation", ""))
+        
+        if enzyme_id and enzyme_id != "nan":
+            available_enzymes[enzyme_id] = {
+                "campaign_id": campaign_id,
+                "has_sequence": bool(protein_seq and protein_seq != "nan"),
+                "generation": generation
+            }
+    
+    identified_count = 0
+    for entry in entries_needing_parents:
+        enzyme_id = entry["enzyme_id"]
+        campaign_id = entry["campaign_id"]
+        generation = entry["generation"]
+        
+        # Create context for Gemini
+        context_info = []
+        context_info.append(f"Enzyme ID: {enzyme_id}")
+        context_info.append(f"Campaign ID: {campaign_id}")
+        if generation:
+            context_info.append(f"Generation: {generation}")
+        
+        # Add available enzymes from the same campaign for context
+        campaign_enzymes = []
+        for enz_id, enz_data in available_enzymes.items():
+            if enz_data["campaign_id"] == campaign_id:
+                status = "with sequence" if enz_data["has_sequence"] else "without sequence"
+                gen_info = f"(gen {enz_data['generation']})" if enz_data["generation"] else ""
+                campaign_enzymes.append(f"  - {enz_id} {status} {gen_info}")
+        
+        if campaign_enzymes:
+            context_info.append("Available enzymes in same campaign:")
+            context_info.extend(campaign_enzymes[:10])  # Limit to first 10 for context
+        
+        context_text = "\n".join(context_info)
+        
+        prompt = f"""
+Based on the enzyme information provided, can you identify the parent enzyme for this enzyme?
+
+{context_text}
+
+This enzyme currently has no sequence data and no parent information. Based on the enzyme ID and the available enzymes in the same campaign, can you identify which enzyme is likely the parent?
+
+Please provide your response in this format:
+Parent: [parent_enzyme_id or "Unknown"]
+
+If you cannot identify a parent enzyme, just respond with "Parent: Unknown".
+"""
+        
+        try:
+            response = model.generate_content(prompt)
+            response_text = response.text.strip()
+            
+            # Parse the response
+            parent_match = re.search(r'Parent:\s*([^\n]+)', response_text)
+            
+            if parent_match:
+                parent = parent_match.group(1).strip()
+                if parent and parent != "Unknown" and parent != "No parent identified":
+                    # Verify the parent exists in our available enzymes
+                    if parent in available_enzymes:
+                        df.at[entry["idx"], "parent_enzyme_id"] = parent
+                        identified_count += 1
+                        log.info(f"Identified parent for {enzyme_id}: {parent}")
+                    else:
+                        log.warning(f"Gemini suggested parent {parent} for {enzyme_id}, but it's not in available enzymes")
+            
+        except Exception as e:
+            log.warning(f"Failed to identify parent for {enzyme_id} from Gemini: {e}")
+            continue
+    
+    if identified_count > 0:
+        log.info(f"Successfully identified {identified_count} parent enzymes using Gemini API")
+    else:
+        log.info("No parent enzymes were identified using Gemini API")
+    
+    return df
+
+
+# === 8. MAIN PROCESSOR === ---------------------------------------------------
 
 class SequenceProcessor:
     """Main processor for handling the complete workflow."""
@@ -866,6 +1005,17 @@ class SequenceProcessor:
             self.process_remaining()
             self.backward_pass()
             
+            # Use Gemini to identify parent enzymes for entries with missing sequences
+            log.info(f"Identifying parents with Gemini for campaign: {campaign_id}")
+            self.df = identify_parents_with_gemini(self.df)
+            
+            # Rebuild relationships after parent identification
+            self.generator = SequenceGenerator(self.df)
+            
+            # Try to fill sequences again after parent identification
+            log.info(f"Attempting to fill sequences after parent identification for campaign: {campaign_id}")
+            self.process_remaining()
+            
             # Update the original dataframe with results
             original_df.loc[campaign_mask, :] = self.df
             

{debase-0.4.1 → debase-0.4.3}/src/debase/enzyme_lineage_extractor.py

@@ -142,21 +142,36 @@ def extract_text(pdf_path: str | Path | bytes) -> str:
 
 
 def extract_captions(pdf_path: str | Path | bytes, max_chars: int = MAX_CHARS) -> str:
-    """Extract figure/table captions using the improved regex.
+    """Extract ALL figure/table captions with extensive surrounding context.
 
     The function scans every text line on every page and keeps lines whose first
     token matches `_CAPTION_PREFIX_RE`. This covers labels such as:
-      * Fig. 1, Figure 2A, Extended Data Fig 3
+      * Fig. 1, Figure 2A, Figure 2B, Figure 2C (ALL sub-captions)
       * Table S1, Table 4, Scheme 2, Chart 1B
-      * Supplementary Fig. S5, Supp Table 2
+      * Supplementary Fig. S5A, S5B, S5C (ALL variations)
+      
+    For SI documents, includes extensive context since understanding what each 
+    section contains is crucial for accurate location identification.
     """
 
     doc = _open_doc(pdf_path)
     captions: list[str] = []
     try:
-        for page in doc:
+        for page_num, page in enumerate(doc):
             page_dict = page.get_text("dict")
+            
+            # Get all text blocks on this page for broader context
+            page_text_blocks = []
             for block in page_dict.get("blocks", []):
+                block_text = ""
+                for line in block.get("lines", []):
+                    text_line = "".join(span["text"] for span in line.get("spans", []))
+                    if text_line.strip():
+                        block_text += text_line.strip() + " "
+                if block_text.strip():
+                    page_text_blocks.append(block_text.strip())
+            
+            for block_idx, block in enumerate(page_dict.get("blocks", [])):
                 # Get all lines in this block
                 block_lines = []
                 for line in block.get("lines", []):
@@ -166,21 +181,94 @@ def extract_captions(pdf_path: str | Path | bytes, max_chars: int = MAX_CHARS) -
                 # Check if any line starts with a caption prefix
                 for i, line in enumerate(block_lines):
                     if _CAPTION_PREFIX_RE.match(line):
-                        # Found a caption start - collect this line and subsequent lines
-                        # until we hit an empty line or the end of the block
+                        context_parts = []
+                        
+                        # Add page context for SI documents (more critical there)
+                        context_parts.append(f"Page {page_num + 1}")
+                        
+                        # Add extensive context before the caption (5-7 lines for SI context)
+                        context_before = []
+                        
+                        # First try to get context from current block
+                        for k in range(max(0, i-7), i):
+                            if k < len(block_lines) and block_lines[k].strip():
+                                if not _CAPTION_PREFIX_RE.match(block_lines[k]):
+                                    context_before.append(block_lines[k])
+                        
+                        # If not enough context, look at previous text blocks on the page
+                        if len(context_before) < 3 and block_idx > 0:
+                            prev_block_text = page_text_blocks[block_idx - 1] if block_idx < len(page_text_blocks) else ""
+                            if prev_block_text:
+                                # Get last few sentences from previous block
+                                sentences = prev_block_text.split('. ')
+                                context_before = sentences[-2:] + context_before if len(sentences) > 1 else [prev_block_text] + context_before
+                        
+                        if context_before:
+                            # Include more extensive context for better understanding
+                            context_text = " ".join(context_before[-5:])  # Last 5 lines/sentences of context
+                            context_parts.append("Context: " + context_text)
+                        
+                        # Extract the COMPLETE caption including all sub-parts
                         caption_parts = [line]
-                        for j in range(i + 1, len(block_lines)):
+                        j = i + 1
+                        
+                        # Continue collecting caption text until we hit a clear break
+                        while j < len(block_lines):
                             next_line = block_lines[j]
-                            if not next_line:  # Empty line signals end of caption
-                                break
-                            # Check if next line is a new caption
+                            
+                            # Stop if we hit an empty line followed by non-caption text
+                            if not next_line:
+                                # Check if the line after empty is a new caption
+                                if j + 1 < len(block_lines) and _CAPTION_PREFIX_RE.match(block_lines[j + 1]):
+                                    break
+                                # If next non-empty line is not a caption, continue collecting
+                                elif j + 1 < len(block_lines):
+                                    j += 1
+                                    continue
+                                else:
+                                    break
+                            
+                            # Stop if we hit a new caption
                             if _CAPTION_PREFIX_RE.match(next_line):
                                 break
+                            
+                            # Include this line as part of the caption
                             caption_parts.append(next_line)
+                            j += 1
                         
-                        # Join the parts with spaces
+                        # Join the caption parts
                         full_caption = " ".join(caption_parts)
-                        captions.append(full_caption)
+                        context_parts.append("Caption: " + full_caption)
+                        
+                        # Add extensive context after the caption (especially important for SI)
+                        context_after = []
+                        
+                        # Look for descriptive text following the caption
+                        for k in range(j, min(len(block_lines), j + 10)):  # Look ahead up to 10 lines
+                            if k < len(block_lines) and block_lines[k].strip():
+                                if not _CAPTION_PREFIX_RE.match(block_lines[k]):
+                                    context_after.append(block_lines[k])
+                        
+                        # If not enough context, look at next text blocks
+                        if len(context_after) < 3 and block_idx + 1 < len(page_text_blocks):
+                            next_block_text = page_text_blocks[block_idx + 1]
+                            if next_block_text:
+                                # Get first few sentences from next block
+                                sentences = next_block_text.split('. ')
+                                context_after.extend(sentences[:3] if len(sentences) > 1 else [next_block_text])
+                        
+                        if context_after:
+                            # Include extensive following context
+                            following_text = " ".join(context_after[:7])  # First 7 lines of following context
+                            context_parts.append("Following: " + following_text)
+                        
+                        # For SI documents, add section context if this appears to be a section header
+                        if any(keyword in full_caption.lower() for keyword in ['supplementary', 'supporting', 'si ', 's1', 's2', 's3']):
+                            context_parts.append("SI_SECTION: This appears to be supplementary material content")
+                        
+                        # Combine all parts with proper separation
+                        full_caption_with_context = " | ".join(context_parts)
+                        captions.append(full_caption_with_context)
     finally:
         doc.close()
 
@@ -645,11 +733,13 @@ find the COMPLETE evolutionary lineage of enzyme variants (i.e. which variant
 came from which parent and what mutations were introduced).
 
 Respond with a JSON array of objects, each containing:
-- "location": the identifier (e.g. "Table S1", "Figure 2B", "p. 6")
+- "location": the identifier (e.g. "Table S1", "Figure 2B", "6" for page 6, "S6" for supplementary page 6)
 - "type": one of "table", "figure", "text", "section"  
 - "confidence": your confidence score (0-100) that this location contains lineage data
 - "reason": brief explanation of why this location likely contains lineage
 
+IMPORTANT: For page numbers, use ONLY the number (e.g., "6" not "p. 6" or "page 6")
+
 Order by confidence score (highest first). Tables showing complete variant lineages or 
 mutation lists should be ranked higher than figure showing complete variant lineages.
 Text sections is used when no suitable tables/figurews exist.
@@ -747,7 +837,7 @@ def identify_campaigns(
     debug_dir: str | Path | None = None,
 ) -> List[Campaign]:
     """Identify distinct directed evolution campaigns in the manuscript."""
-    prompt = _CAMPAIGN_IDENTIFICATION_PROMPT.format(text=text[:30_000])
+    prompt = _CAMPAIGN_IDENTIFICATION_PROMPT.format(text=text)
     campaigns_data: List[dict] = []
     try:
         campaigns_data = generate_json_with_retry(
@@ -825,7 +915,7 @@ def identify_evolution_locations(
     
     # Include TOC before the main text
     combined_text = toc_text + text if toc_text else text
-    prompt = _LINEAGE_LOC_PROMPT.format(max_results=max_results) + "\n\nTEXT:\n" + combined_text[:15_000]
+    prompt = _LINEAGE_LOC_PROMPT.format(max_results=max_results) + "\n\nTEXT:\n" + combined_text
     locs: List[dict] = []
     try:
         locs = generate_json_with_retry(
@@ -1306,7 +1396,7 @@ def get_lineage(
     5. Return both variants and campaigns.
     """
     # First, identify campaigns in the manuscript
-    campaigns = identify_campaigns(full_text[:50_000], model, debug_dir=debug_dir)
+    campaigns = identify_campaigns(full_text, model, debug_dir=debug_dir)
     
     if campaigns:
         log.info(f"Identified {len(campaigns)} distinct campaigns")
@@ -1364,7 +1454,7 @@ def get_lineage(
             context_text = _extract_text_at_locations(full_text, [loc], context_chars=500)
             locations_with_context.append({
                 'location': loc,
-                'context': context_text[:1000]  # First 1000 chars of extracted context
+                'context': context_text  # Full extracted context
             })
         
         # For each campaign, ask Gemini to select the best location
@@ -1554,13 +1644,17 @@ PRIORITY: Protein/amino acid sequences are preferred over DNA sequences.
 Look for table of contents entries or section listings that mention sequences.
 Return a JSON array where each element has:
 - "section": the section heading or description
-- "page": the page number shown in the table of contents for this section, to your best judgement.
+- "page": the page number (IMPORTANT: Return ONLY the number, e.g., "53" not "p. 53" or "page 53")
 
 Focus on:
 - Table of contents or entries about "Sequence Information" or "Nucleotide and amino acid sequences"
-- Return the EXACT notation as shown.
+- For supplementary pages, use "S" prefix (e.g., "S53" not "p. S53")
 - Prioritize sections that mention "protein" or "amino acid" sequences
 
+CRITICAL: Page numbers must be returned as plain numbers or S-prefixed numbers only:
+- Correct: "53", "S12", "147"
+- Wrong: "p. 53", "P. 53", "page 53", "pg 53"
+
 Return [] if no sequence sections are found.
 Absolutely don't include nucleotides or primer sequences, it is better to return nothing then incomplete sequence, use your best judgement.
 
@@ -1572,7 +1666,7 @@ TEXT (truncated):
 
 def identify_sequence_locations(text: str, model, *, debug_dir: str | Path | None = None) -> list[dict]:
     """Ask Gemini for promising places to look for sequences."""
-    prompt = _SEQ_LOC_PROMPT.format(chunk=text[:15_000])
+    prompt = _SEQ_LOC_PROMPT.format(chunk=text)
     try:
         locs = generate_json_with_retry(model, prompt, debug_dir=debug_dir, tag="seq_locations")
         return locs if isinstance(locs, list) else []