diff --git a/.gitignore b/.gitignore index 2d6a3de..9444d3f 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ doc/ build/ dist/ grobid_client_python.egg-info/ +.venv \ No newline at end of file diff --git a/Readme.md b/Readme.md index 31f6f19..7b98490 100644 --- a/Readme.md +++ b/Readme.md @@ -144,6 +144,9 @@ grobid_client [OPTIONS] SERVICE | `--flavor` | Processing flavor for fulltext extraction | | `--json` | Convert TEI output to JSON format | | `--markdown` | Convert TEI output to Markdown format | +| `--typed_area` | Enable sending typed-area layout JSON to GROBID | +| `--typed_areas_dir` | Directory of pre-computed JSON files | +| `--typed_area_server` | URL of the PaddlePaddle server | #### Examples @@ -166,8 +169,22 @@ grobid_client --server https://grobid.example.com --input ~/citations.txt proces # Force reprocessing with sentence segmentation and JSON output grobid_client --input ~/docs --force --segment_sentences --json processFulltextDocument + +# Typed Area Processing (with PaddlePaddle server) +# The client will fetch JSON from the paddle server and send it to Grobid +grobid_client --input ~/pdfs --output ~/results --typed_area --typed_area_server h processFulltextDocument + +# Typed Area Processing (with pre-computed offline JSON files) +grobid_client --input ~/pdfs --output ~/results --typed_area --typed_areas_dir ~/precomputed_jsons processFulltextDocument ``` +### Worker and Concurrency (Typed Areas) +When using the `--typed_area_server` flag, the Grobid client makes requests to *both* the PaddlePaddle server and the Grobid server. + +- **Grobid Client Threads (`--n`)**: Controls how many PDFs are processed concurrently. +- **PaddlePaddle Server Workers (`--workers`)**: Controls how many concurrent layout requests the PaddlePaddle server can handle. + +**Recommendation:** Set the PaddlePaddle server `--workers` to match the grobid-client `--n` threads (e.g., `--n 4` on the client and `--workers 4` on the server). ### Python Library #### Basic Usage diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index a525b9f..4bbafb7 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -24,7 +24,7 @@ import requests import pathlib import logging -from typing import Tuple +from typing import Any, Optional, Tuple, Union, List import copy from .format.TEI2LossyJSON import TEI2LossyJSONConverter @@ -34,7 +34,7 @@ class ServerUnavailableException(Exception): """Exception raised when GROBID server is not available or not responding.""" - def __init__(self, message="GROBID server is not available"): + def __init__(self, message: str = "GROBID server is not available") -> None: super().__init__(message) self.message = message @@ -77,15 +77,15 @@ class GrobidClient(ApiClient): def __init__( self, - grobid_server=None, - batch_size=None, - coordinates=None, - sleep_time=None, - timeout=None, - config_path=None, - check_server=True, - verbose=False - ): + grobid_server: Optional[str] = None, + batch_size: Optional[int] = None, + coordinates: Optional[List[str]] = None, + sleep_time: Optional[int] = None, + timeout: Optional[int] = None, + config_path: Optional[str] = None, + check_server: bool = True, + verbose: bool = False + ) -> None: # Store verbose parameter for logging configuration self.verbose = verbose @@ -112,13 +112,13 @@ def __init__( if check_server: self._test_server_connection() - def _set_config_params(self, params): + def _set_config_params(self, params: dict) -> None: """Set configuration parameters, only if they are not None.""" for key, value in params.items(): if value is not None: self.config[key] = value - def _warn_on_consolidation_timeout(self, consolidate_citations): + def _warn_on_consolidation_timeout(self, consolidate_citations: bool) -> None: """Warn when citation consolidation is enabled with a low client timeout. Consolidating citations makes GROBID query external services and can be @@ -139,23 +139,25 @@ def _warn_on_consolidation_timeout(self, consolidate_citations): f"(2-3 minutes is recommended)." ) - def _handle_server_busy_retry(self, file_path, retry_func, *args, **kwargs): + def _handle_server_busy_retry(self, file_path: str, retry_func: Any, *args: Any, **kwargs: Any) -> Any: """Handle server busy (503) retry logic.""" self.logger.warning(f"Server busy (503), retrying {file_path} after {self.config['sleep_time']} seconds") time.sleep(self.config["sleep_time"]) return retry_func(*args, **kwargs) - def _handle_request_error(self, file_path, error, error_type="Request"): + def _handle_request_error( + self, file_path: str, error: Exception, error_type: str = "Request" + ) -> Tuple[str, int, str]: """Handle request errors with consistent logging and return format.""" self.logger.error(f"{error_type} failed for {file_path}: {str(error)}") return (file_path, 500, f"{error_type} failed: {str(error)}") - def _handle_unexpected_error(self, file_path, error): + def _handle_unexpected_error(self, file_path: str, error: Exception) -> Tuple[str, int, str]: """Handle unexpected errors with consistent logging and return format.""" self.logger.error(f"Unexpected error processing {file_path}: {str(error)}") return (file_path, 500, f"Unexpected error: {str(error)}") - def _configure_logging(self): + def _configure_logging(self) -> None: """Configure logging based on the configuration settings.""" # Get logging config with defaults log_config = self.config.get('logging', {}) @@ -231,7 +233,7 @@ def _configure_logging(self): self.logger.info( f"Logging configured - Level: {log_level_str}, Console: {log_config.get('console', True)}, File: {log_file or 'disabled'}") - def _parse_file_size(self, size_str): + def _parse_file_size(self, size_str: Union[str, int]) -> int: """Parse file size string like '10MB', '1GB' to bytes.""" size_str = str(size_str).upper().strip() @@ -255,7 +257,7 @@ def _parse_file_size(self, size_str): return int(number * multipliers.get(unit, 1)) - def _load_config(self, path="./config.json"): + def _load_config(self, path: str = "./config.json") -> None: """ Load and merge configuration from a JSON file with default values. If the file doesn't exist, keep the default values. @@ -327,7 +329,12 @@ def _test_server_connection(self) -> Tuple[bool, int]: self.logger.error(error_msg) raise ServerUnavailableException(error_msg) from e - def _output_file_name(self, input_file, input_path, output): + def _output_file_name( + self, + input_file: str, + input_path: str, + output: Optional[str], + ) -> str: # Use pathlib for consistent cross-platform path handling input_file_path = pathlib.Path(input_file) @@ -351,23 +358,26 @@ def ping(self) -> Tuple[bool, int]: def process( self, - service, - input_path, - output=None, - n=10, - generate_ids=False, - consolidate_header=True, - consolidate_citations=False, - include_raw_citations=False, - include_raw_affiliations=False, - tei_coordinates=False, - segment_sentences=False, - force=True, - verbose=False, - flavor=None, - json_output=False, - markdown_output=False - ): + service: str, + input_path: str, + output: Optional[str] = None, + n: int = 10, + generate_ids: bool = False, + consolidate_header: bool = True, + consolidate_citations: bool = False, + include_raw_citations: bool = False, + include_raw_affiliations: bool = False, + tei_coordinates: bool = False, + segment_sentences: bool = False, + force: bool = True, + verbose: bool = False, + flavor: Optional[str] = None, + json_output: bool = False, + markdown_output: bool = False, + typed_area: bool = False, + typed_areas_dir: Optional[str] = None, + typed_area_server: Optional[str] = None + ) -> None: start_time = time.time() batch_size_pdf = self.config["batch_size"] @@ -435,7 +445,10 @@ def process( verbose, flavor, json_output, - markdown_output + markdown_output, + typed_area, + typed_areas_dir, + typed_area_server ) processed_files_count += batch_processed errors_files_count += batch_errors @@ -461,7 +474,10 @@ def process( verbose, flavor, json_output, - markdown_output + markdown_output, + typed_area, + typed_areas_dir, + typed_area_server ) processed_files_count += batch_processed errors_files_count += batch_errors @@ -483,24 +499,27 @@ def process( def process_batch( self, - service, - input_files, - input_path, - output, - n, - generate_ids, - consolidate_header, - consolidate_citations, - include_raw_citations, - include_raw_affiliations, - tei_coordinates, - segment_sentences, - force, - verbose=False, - flavor=None, - json_output=False, - markdown_output=False - ): + service: str, + input_files: List[str], + input_path: str, + output: Optional[str], + n: int, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + force: bool, + verbose: bool = False, + flavor: Optional[str] = None, + json_output: bool = False, + markdown_output: bool = False, + typed_area: bool = False, + typed_areas_dir: Optional[str] = None, + typed_area_server: Optional[str] = None + ) -> Tuple[int, int, int]: batch_start_time = time.time() if verbose: self.logger.info(f"{len(input_files)} files to process in current batch") @@ -584,7 +603,10 @@ def process_batch( segment_sentences, flavor, -1, - -1) + -1, + typed_area, + typed_areas_dir, + typed_area_server) results.append(r) @@ -668,21 +690,93 @@ def process_batch( return processed_count, error_count, skipped_count + def _resolve_typed_area(self, pdf_file: str, typed_areas_dir: Optional[str], typed_area_server: Optional[str]) -> Optional[str]: + """Resolve typed-area JSON for a PDF file. + + Priority: + 1. PaddlePaddle server (if typed_area_server is set) + 2. Explicit directory (if typed_areas_dir is set) + 3. Same directory as the PDF + + Returns: + str or None: JSON string to attach as the typedAreas data field, + or None if no typed-area data could be resolved. + """ + stem = pathlib.Path(pdf_file).stem + + #query the PaddlePaddle server + if typed_area_server: + try: + with open(pdf_file, "rb") as f: + resp = requests.post( + f"{typed_area_server.rstrip('/')}", + files={"file": (os.path.basename(pdf_file), f, "application/pdf")}, + timeout=self.config["timeout"] + ) + if resp.status_code == 200: + json_data = resp.json() + + # Save the JSON to disk + save_dir = typed_areas_dir if typed_areas_dir else str(pathlib.Path(pdf_file).parent) + os.makedirs(save_dir, exist_ok=True) + json_path = os.path.join(save_dir, f"{stem}.json") + try: + with open(json_path, "w", encoding="utf-8") as f: + json.dump(json_data, f, ensure_ascii=False, indent=2) + self.logger.debug(f"Saved typed-area JSON to {json_path}") + except Exception as e: + self.logger.warning(f"Failed to save typed-area JSON to {json_path}: {e}") + + # Extract just the elements array if present, as GROBID expects a JSON Array + payload = json_data.get("elements", json_data) if isinstance(json_data, dict) else json_data + return json.dumps(payload) + else: + self.logger.warning( + f"Typed-area server returned {resp.status_code} for {pdf_file}" + ) + except Exception as e: + self.logger.warning( + f"Typed-area server request failed for {pdf_file}: {e}" + ) + return None + + # explicit directory, or same directory as the PDF + search_dir = typed_areas_dir if typed_areas_dir else str(pathlib.Path(pdf_file).parent) + json_path = os.path.join(search_dir, f"{stem}.json") + + if os.path.isfile(json_path): + try: + with open(json_path, "r", encoding="utf-8") as f: + json_content = json.load(f) + payload = json_content.get("elements", json_content) if isinstance(json_content, dict) else json_content + return json.dumps(payload) + except Exception as e: + self.logger.warning(f"Failed to read typed-area JSON {json_path}: {e}") + return None + + self.logger.warning( + f"No typed-area JSON found for {pdf_file} (looked in {search_dir})" + ) + return None + def process_pdf( self, - service, - pdf_file, - generate_ids, - consolidate_header, - consolidate_citations, - include_raw_citations, - include_raw_affiliations, - tei_coordinates, - segment_sentences, - flavor=None, - start=-1, - end=-1 - ): + service: str, + pdf_file: str, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + flavor: Optional[str] = None, + start: int = -1, + end: int = -1, + typed_area: bool = False, + typed_areas_dir: Optional[str] = None, + typed_area_server: Optional[str] = None + ) -> Tuple[str, int, Optional[str]]: pdf_handle = None try: pdf_handle = open(pdf_file, "rb") @@ -721,6 +815,14 @@ def process_pdf( if end and end > 0: the_data["end"] = str(end) + # Resolve and attach typed-area JSON if enabled + if typed_area: + typed_area_json = self._resolve_typed_area( + pdf_file, typed_areas_dir, typed_area_server + ) + if typed_area_json: + the_data["typedAreas"] = typed_area_json + res, status = self.post( url=the_url, files=files, data=the_data, headers={"Accept": "text/plain"}, timeout=self.config['timeout'] @@ -741,7 +843,10 @@ def process_pdf( segment_sentences, flavor, start, - end + end, + typed_area, + typed_areas_dir, + typed_area_server ) return (pdf_file, status, res.text) @@ -760,24 +865,24 @@ def process_pdf( if pdf_handle: pdf_handle.close() - def get_server_url(self, service): + def get_server_url(self, service: str) -> str: return self.config['grobid_server'] + "/api/" + service def process_txt( self, - service, - txt_file, - generate_ids, - consolidate_header, - consolidate_citations, - include_raw_citations, - include_raw_affiliations, - tei_coordinates, - segment_sentences, - flavor=None, - start_page=-1, - end_page=-1 - ): + service: str, + txt_file: str, + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + flavor: Optional[str] = None, + start_page: int = -1, + end_page: int = -1 + ) -> Tuple[str, int, Optional[str]]: # create request based on file content try: with open(txt_file, 'r', encoding='utf-8') as f: @@ -792,7 +897,7 @@ def process_txt( the_url = self.get_server_url(service) # set the GROBID parameters - the_data = {} + the_data: dict = {} if consolidate_citations: the_data["consolidateCitations"] = "1" if include_raw_citations: @@ -826,7 +931,7 @@ def process_txt( return (txt_file, status, res.text) -def main(): +def main() -> None: # Basic logging setup for initialization only # The actual logging configuration will be done by GrobidClient based on config.json temp_logger = logging.getLogger(__name__) @@ -941,6 +1046,21 @@ def main(): action="store_true", help="Convert TEI output to Markdown format", ) + parser.add_argument( + "--typed_area", + action="store_true", + help="Enable typed-area support: attach PaddlePaddle layout JSON to each Grobid request", + ) + parser.add_argument( + "--typed_areas_dir", + default=None, + help="Directory containing pre-computed typed-area JSON files (default: same directory as the PDF)", + ) + parser.add_argument( + "--typed_area_server", + default=None, + help="URL of the PaddlePaddle typed-area server", + ) args = parser.parse_args() @@ -950,6 +1070,9 @@ def main(): flavor = args.flavor json_output = args.json markdown_output = args.markdown + typed_area = args.typed_area + typed_areas_dir = args.typed_areas_dir + typed_area_server = args.typed_area_server # Initialize n with default value n = 10 @@ -1020,7 +1143,10 @@ def main(): verbose=verbose, flavor=flavor, json_output=json_output, - markdown_output=markdown_output + markdown_output=markdown_output, + typed_area=typed_area, + typed_areas_dir=typed_areas_dir, + typed_area_server=typed_area_server ) except Exception as e: logger.error(f"Processing failed: {str(e)}")