Merge pull request #22 from SFBioinformaticsGroup/hybrid-mode

Hybrid mode

Author: flbulgarelli

setup.cfg

@@ -54,6 +54,10 @@ install_requires =
     tqdm >= 4.67.1
     camelot-py[cv] == 1.0.9
     pdfplumber >= 0.11
+    img2table >= 1.4.2
+    opencv-contrib-python >= 4.12.0
+    PyMuPDF >= 1.26
+    pymupdf-layout >= 1.26
     pandas == 2.3.2
 
 [options.packages.find]

src/paper2table/__main__.py

@@ -3,11 +3,13 @@
 import sys
 import time
 from pathlib import Path
+from typing import Optional
 
 from tqdm import tqdm
 
 from paper2table import __version__
-from paper2table.readers import agent, camelot, pdfplumber
+from paper2table.mapping import TablesMapping
+from paper2table.readers import agent, camelot, pdfplumber, hybrid
 from paper2table.tables_reader import TablesReader
 from paper2table.writers import file, stdout, tablemerge
 from paper2table.writers.tablemerge import TablemergeMetadata
@@ -42,6 +44,14 @@ def parse_args():
         help="How tables are going to be extracted",
         default="pdfplumber",
     )
+    parser.add_argument(
+        "-H",
+        "--hybrid",
+        dest="hybrid",
+        help="Enable hybrid mode",
+        action="store_const",
+        const=True,
+    )
     parser.add_argument(
         "-m",
         "--model",
@@ -61,19 +71,25 @@ def parse_args():
         "-s",
         "--schema",
         type=str,
-        help="set table schema in the form column:type. Only used by agent reader",
+        help="set table schema in the form column:type. Only used by agent or hybrid reader",
     )
     parser.add_argument(
         "-p",
         "--schema-path",
         type=str,
         help="set table schema path. Only used by agent reader",
     )
+    parser.add_argument(
+        "-M",
+        "--mappings-path",
+        type=str,
+        help="set tables mapping path. Only used by hybrid reader",
+    )
     parser.add_argument(
         "-c",
         "--column-names-hints-path",
         type=str,
-        help="set table schema path. Only used by agent reader",
+        help="set table schema path. Only used by pdfplumber reader",
     )
     parser.add_argument(
         "-o",
@@ -112,44 +128,79 @@ def setup_logging(loglevel):
     """
     logformat = "[%(asctime)s] %(levelname)s:%(name)s:%(message)s"
     logging.basicConfig(
-        level=loglevel, stream=sys.stdout, format=logformat, datefmt="%Y-%m-%d %H:%M:%S"
+        stream=sys.stdout, format=logformat, datefmt="%Y-%m-%d %H:%M:%S"
     )
+    logging.getLogger().setLevel(logging.WARN)
+    if loglevel:
+        _logger.setLevel(loglevel)
 
 
 def get_tables_reader(args):
     if args.reader == "agent":
         schema = Path(args.schema_path).read_text() if args.schema_path else args.schema
         if not schema:
-            print("Missing schema. Need to either pass --schema-path or --schema")
+            print(
+                "Missing schema. Need to either pass --schema-path or --schema when using agent reader"
+            )
             sys.exit(1)
 
-        def read_tables(paper_path: str):
+        def read_tables(paper_path: str, _mapping: Optional[TablesMapping] = None):
             time.sleep(args.model_sleep)
-            _logger.debug(
-                f"Processing paper {paper_path} with model {args.model} and {schema}..."
-            )
+            _logger.debug(f"Processing paper {paper_path} with model {args.model}")
             return agent.read_tables(paper_path, model=args.model, schema=schema)
 
     elif args.reader == "pdfplumber":
+        column_names_hints = (
+            Path(args.column_names_hints_path).read_text()
+            if args.column_names_hints_path
+            else ""
+        )
 
-        def read_tables(paper_path: str):
-            column_names_hints = (
-                Path(args.column_names_hints_path).read_text()
-                if args.column_names_hints_path
-                else ""
-            )
+        _logger.debug(
+            f"Using pdfplumber reader with column names hints {column_names_hints}"
+        )
+
+        def read_tables(paper_path: str, mapping: Optional[TablesMapping] = None):
 
-            _logger.debug(
-                f"Processing paper {paper_path} with pdfplumber and {column_names_hints} as column names hints..."
+            _logger.debug(f"Processing paper {paper_path}...")
+            return pdfplumber.read_tables(
+                paper_path, column_names_hints, mapping=mapping
             )
-            return pdfplumber.read_tables(paper_path, column_names_hints)
 
     else:
+        _logger.debug(f"Using camelot reader {args.reader}-{args.model}")
 
-        def read_tables(paper_path: str):
-            _logger.debug(f"Processing paper {paper_path} with camelot...")
+        def read_tables(paper_path: str, _mapping: Optional[TablesMapping] = None):
+            _logger.debug(f"Processing paper {paper_path}...")
             return camelot.read_tables(paper_path)
 
+    if args.hybrid:
+        mappings_path = (
+            Path(args.mappings_path) if args.schema_path else Path("./mappings")
+        )
+        schema = Path(args.schema_path).read_text() if args.schema_path else args.schema
+        if not schema:
+            print(
+                "Missing schema. Need to either pass --schema-path or --schema when using hybrid mode"
+            )
+            sys.exit(1)
+
+        _logger.debug(f"Schema is {schema}")
+        _logger.debug(f"Applying {args.reader}-{args.model} hybrid reader")
+
+        base_reader = read_tables
+
+        def read_tables(paper_path: str, _mapping: Optional[TablesMapping] = None):
+            time.sleep(args.model_sleep)
+            _logger.debug(f"Processing paper {paper_path}")
+            return hybrid.read_tables(
+                paper_path,
+                model=args.model,
+                mappings_path=mappings_path,
+                schema=schema,
+                reader=base_reader,
+            )
+
     return read_tables
 
 

src/paper2table/hints.py

@@ -0,0 +1,6 @@
+
+from utils.normalize_name import normalize_name
+from utils.columns_schema import tokenize_schema
+
+def parse_column_names_hints(hints: str) -> list[str]:
+    return [normalize_name(hint) for hint in tokenize_schema(hints)]

src/paper2table/mapping.py

@@ -0,0 +1,48 @@
+
+
+from typing import Literal
+from pydantic import BaseModel
+
+
+class ColumnMapping(BaseModel):
+    from_column_number: int
+    """
+    The original column number
+    """
+
+    to_column_name: str
+    """
+    The desired column name
+    """
+
+
+class TableMapping(BaseModel):
+    """
+    Instructions for read_table
+    about how to read a table.
+    """
+
+    title: str
+
+    header_mode: Literal["all_pages", "first_page_only", "none"]
+
+    first_page: int
+    """
+    1-based first page number where table is allocated
+    """
+
+    last_page: int
+    """
+    1-based last page number where table is allocated
+    """
+
+    column_mappings: list[ColumnMapping]
+    """
+    Mappings that go from original column number
+    to desired column name
+    """
+
+
+class TablesMapping(BaseModel):
+    tables: list[TableMapping]
+    citation: str

src/paper2table/readers/agent.py

@@ -1,37 +1,13 @@
 from pathlib import Path
-from typing import Any
 
 from pydantic import create_model
 from pydantic_ai import Agent, BinaryContent
 
-from utils.tokenize_schema import tokenize_schema
+from utils.columns_schema import parse_schema
 
 from ..tables_reader import TablesReader
 from ..tables_reader.pydantic import TablesModelWrapper
 
-types_map: dict[str, Any] = {
-    "str": str,
-    "int": int,
-    "float": float,
-    "bool": bool,
-}
-
-
-def parse_schema(schema_str: str) -> dict[str, tuple[Any, ...]]:
-    parts = tokenize_schema(schema_str)
-
-    fields: dict[str, tuple[Any, ...]] = {}
-    for part in parts:
-        if ":" not in part:
-            raise ValueError(f"Invalid field specifier: {part}")
-        name, type_str = part.split(":", 1)
-        if type_str not in types_map:
-            raise ValueError(f"Unsupported type: {type_str}")
-        fields[name] = (types_map[type_str], ...)
-
-    return fields
-
-
 def build_table_model(schema: str):
     """
     Build and return a TableModel from a schema string.