debase 0.1.4__tar.gz → 0.1.6__tar.gz

{debase-0.1.4 → debase-0.1.6}/PKG-INFO

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

{debase-0.1.4 → debase-0.1.6}/src/debase/_version.py

@@ -1,3 +1,3 @@
 """Version information."""
 
-__version__ = "0.1.4"
+__version__ = "0.1.6"

{debase-0.1.4 → debase-0.1.6}/src/debase/enzyme_lineage_extractor.py

@@ -1024,6 +1024,38 @@ IMPORTANT: Only extract variants that belong to this specific campaign.
 
 # ---- 6.3  Helper for location-based extraction -----------------------------
 
+def _is_toc_entry(text: str, position: int, pattern: str) -> bool:
+    """Check if a found pattern is likely a table of contents entry."""
+    # Find the line containing this position
+    line_start = text.rfind('\n', 0, position)
+    line_end = text.find('\n', position)
+    
+    if line_start == -1:
+        line_start = 0
+    else:
+        line_start += 1
+        
+    if line_end == -1:
+        line_end = len(text)
+        
+    line = text[line_start:line_end]
+    
+    # TOC indicators:
+    # 1. Line contains dots (...) followed by page number
+    # 2. Line ends with just a page number
+    # 3. Line has "Table S12:" or similar followed by title and page
+    if '...' in line or re.search(r'\.\s*\d+\s*$', line) or re.search(r':\s*[^:]+\s+\d+\s*$', line):
+        return True
+        
+    # Check if this is in a contents/TOC section
+    # Look backwards up to 500 chars for "Contents" or "Table of Contents"
+    context_start = max(0, position - 500)
+    context = text[context_start:position].lower()
+    if 'contents' in context or 'table of contents' in context:
+        return True
+        
+    return False
+
 def _extract_text_at_locations(text: str, locations: List[Union[str, dict]], context_chars: int = 5000, validate_sequences: bool = False) -> str:
     """Extract text around identified locations."""
     if not locations:
@@ -1082,11 +1114,25 @@ def _extract_text_at_locations(text: str, locations: List[Union[str, dict]], con
         pos = -1
         used_pattern = None
         for pattern in page_patterns:
-            temp_pos = text_lower.find(pattern.lower())
-            if temp_pos != -1:
+            search_pos = 0
+            while search_pos < len(text_lower):
+                temp_pos = text_lower.find(pattern.lower(), search_pos)
+                if temp_pos == -1:
+                    break
+                    
+                # Check if this is a TOC entry
+                if _is_toc_entry(text, temp_pos, pattern):
+                    log.debug("Skipping TOC entry for pattern '%s' at position %d", pattern, temp_pos)
+                    search_pos = temp_pos + len(pattern)
+                    continue
+                    
+                # Found non-TOC entry
                 pos = temp_pos
                 used_pattern = pattern
-                log.debug("Found pattern '%s' at position %d", pattern, pos)
+                log.debug("Found pattern '%s' at position %d (not TOC)", pattern, pos)
+                break
+                
+            if pos != -1:
                 break
         
         if pos != -1:
@@ -1275,7 +1321,9 @@ def get_lineage(
                 
             # Use text-based extraction (works for tables and text sections)
             # Extract from full text, not caption text - use only primary location
-            focused_text = _extract_text_at_locations(full_text, [primary_location])
+            # Use more context for tables since they often span multiple pages
+            context_size = 15000 if location_type == 'table' else 5000
+            focused_text = _extract_text_at_locations(full_text, [primary_location], context_chars=context_size)
             log.info("Reduced text from %d to %d chars using primary location %s for campaign %s", 
                      len(full_text), len(focused_text), 
                      primary_location.get('location', 'Unknown') if isinstance(primary_location, dict) else 'Unknown',
@@ -1864,7 +1912,7 @@ def _infer_generations(variants: List[Variant]) -> None:
 
 
 def _merge_lineage_and_sequences(
-    lineage: List[Variant], seqs: List[SequenceBlock], doi: Optional[str]
+    lineage: List[Variant], seqs: List[SequenceBlock], doi: Optional[str], model=None
 ) -> pd.DataFrame:
     """Return a tidy DataFrame with one row per variant."""
 
@@ -1895,6 +1943,65 @@ def _merge_lineage_and_sequences(
     # 2. Outer merge keeps every lineage entry and adds sequence cols when present
     df = pd.merge(df_lin, df_seq, on="variant_id", how="left")
 
+    # 2a. If we have unmatched sequences and a model, use Gemini to match them
+    if model and len(df_seq) > 0:
+        # Find lineage entries without sequences
+        missing_seq = df['aa_seq'].isna() & df['dna_seq'].isna()
+        unmatched_lineage = df[missing_seq]['variant_id'].tolist()
+        
+        # Find sequences that weren't matched
+        matched_seq_ids = df[~missing_seq]['variant_id'].tolist()
+        unmatched_seqs = df_seq[~df_seq['variant_id'].isin(matched_seq_ids)]
+        
+        if len(unmatched_lineage) > 0 and len(unmatched_seqs) > 0:
+            log.info(f"Found {len(unmatched_lineage)} lineage entries without sequences and {len(unmatched_seqs)} unmatched sequences")
+            log.info(f"Using Gemini to match variants")
+            
+            # Build prompt for Gemini to match variants
+            prompt = f"""Match enzyme variant IDs between two lists. The same variant may be labeled differently in different parts of the paper.
+
+Lineage variant IDs (need sequences):
+{json.dumps(unmatched_lineage)}
+
+Sequence variant IDs (have sequences):
+{json.dumps(unmatched_seqs['variant_id'].tolist())}
+
+These lists contain variant identifiers from the same paper but may use different naming conventions.
+Match each lineage variant ID to its corresponding sequence variant ID based on any patterns or relationships you can identify.
+
+Return ONLY a JSON object mapping lineage variant IDs to sequence variant IDs. Include only matches you are confident about.
+"""
+            
+            try:
+                response = model.generate_content(prompt)
+                text = _extract_text(response).strip()
+                
+                # Parse JSON response
+                if text.startswith("```"):
+                    text = text.split("```")[1].strip()
+                    if text.startswith("json"):
+                        text = text[4:].strip()
+                
+                matches = json.loads(text)
+                
+                # Apply the matches
+                for lineage_id, seq_id in matches.items():
+                    if lineage_id in unmatched_lineage and seq_id in unmatched_seqs['variant_id'].values:
+                        # Get the sequence data
+                        seq_data = unmatched_seqs[unmatched_seqs['variant_id'] == seq_id].iloc[0]
+                        
+                        # Update the dataframe
+                        mask = df['variant_id'] == lineage_id
+                        df.loc[mask, 'aa_seq'] = seq_data['aa_seq']
+                        df.loc[mask, 'dna_seq'] = seq_data['dna_seq']
+                        df.loc[mask, 'seq_confidence'] = seq_data['seq_confidence']
+                        df.loc[mask, 'truncated'] = seq_data['truncated']
+                        
+                        log.info(f"Matched {lineage_id} -> {seq_id} using Gemini")
+                
+            except Exception as e:
+                log.warning(f"Failed to match variants using Gemini: {e}")
+
     # 3. If generation missing after user input, try inference
     if df["generation"].isna().any():
         _infer_generations(lineage)  # mutates in place
@@ -1920,6 +2027,7 @@ def merge_and_score(
     lineage: List[Variant],
     seqs: List[SequenceBlock],
     doi: Optional[str] = None,
+    model=None,
 ) -> pd.DataFrame:
     """User-facing helper imported by the pipeline orchestrator.
 
@@ -1932,7 +2040,7 @@ def merge_and_score(
         raise ValueError("merge_and_score(): `lineage` list is empty; nothing to merge")
 
     # If no sequences found, still build a DataFrame so caller can decide what to do.
-    df = _merge_lineage_and_sequences(lineage, seqs, doi)
+    df = _merge_lineage_and_sequences(lineage, seqs, doi, model)
 
     # Basic sanity: warn if many missing sequences
     missing_rate = df["aa_seq"].isna().mean() if "aa_seq" in df else 1.0
@@ -2115,7 +2223,7 @@ def run_pipeline(
 
     # 5. Merge & score (Section 8) --------------------------------------------
     doi = extract_doi(manuscript)
-    df_final = merge_and_score(lineage, sequences, doi)
+    df_final = merge_and_score(lineage, sequences, doi, model)
 
     # 6. Write FINAL CSV -------------------------------------------------------
     if output_csv:

{debase-0.1.4 → debase-0.1.6}/src/debase.egg-info/PKG-INFO

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