wehub-resource-sync
661 lines
22 KiB
Python
661 lines
22 KiB
Python
"""
|
|
Batch and Parallel Document Parsing
|
|
|
|
This module provides functionality for processing multiple documents in parallel,
|
|
with progress reporting and error handling.
|
|
"""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional, Tuple
|
|
from dataclasses import dataclass, field
|
|
import time
|
|
|
|
from tqdm import tqdm
|
|
|
|
from .parser import get_parser
|
|
|
|
|
|
@dataclass
|
|
class BatchProcessingResult:
|
|
"""Result of batch processing operation"""
|
|
|
|
successful_files: List[str]
|
|
failed_files: List[str]
|
|
total_files: int
|
|
processing_time: float
|
|
errors: Dict[str, str]
|
|
output_dir: str
|
|
dry_run: bool = False
|
|
skipped_files: List[str] = field(default_factory=list)
|
|
|
|
@property
|
|
def success_rate(self) -> float:
|
|
"""Calculate success rate as percentage"""
|
|
if self.total_files == 0:
|
|
return 0.0
|
|
completed_files = len(self.successful_files) + len(self.skipped_files)
|
|
return (completed_files / self.total_files) * 100
|
|
|
|
def summary(self) -> str:
|
|
"""Generate a summary of the batch processing results"""
|
|
return (
|
|
f"Batch Processing Summary:\n"
|
|
f" Total files: {self.total_files}\n"
|
|
f" Successful: {len(self.successful_files)}\n"
|
|
f" Failed: {len(self.failed_files)}\n"
|
|
f" Skipped: {len(self.skipped_files)}\n"
|
|
f" Success rate: {self.success_rate:.1f}%\n"
|
|
f" Processing time: {self.processing_time:.2f} seconds\n"
|
|
f" Output directory: {self.output_dir}\n"
|
|
f" Dry run: {self.dry_run}"
|
|
)
|
|
|
|
|
|
class BatchParser:
|
|
"""
|
|
Batch document parser with parallel processing capabilities
|
|
|
|
Supports processing multiple documents concurrently with progress tracking
|
|
and comprehensive error handling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
parser_type: str = "mineru",
|
|
max_workers: int = 4,
|
|
show_progress: bool = True,
|
|
timeout_per_file: int = 300,
|
|
skip_installation_check: bool = False,
|
|
):
|
|
"""
|
|
Initialize batch parser
|
|
|
|
Args:
|
|
parser_type: Type of parser to use ("mineru", "docling", or "paddleocr")
|
|
max_workers: Maximum number of parallel workers
|
|
show_progress: Whether to show progress bars
|
|
timeout_per_file: Timeout in seconds for each file
|
|
skip_installation_check: Skip parser installation check (useful for testing)
|
|
"""
|
|
self.parser_type = parser_type
|
|
self.max_workers = max_workers
|
|
self.show_progress = show_progress
|
|
self.timeout_per_file = timeout_per_file
|
|
self.logger = logging.getLogger(__name__)
|
|
|
|
# Initialize parser
|
|
try:
|
|
self.parser = get_parser(parser_type)
|
|
except ValueError as exc:
|
|
raise ValueError(f"Unsupported parser type: {parser_type}") from exc
|
|
|
|
# Check parser installation (optional)
|
|
if not skip_installation_check:
|
|
if not self.parser.check_installation():
|
|
self.logger.warning(
|
|
f"{parser_type.title()} parser installation check failed. "
|
|
f"This may be due to package conflicts. "
|
|
f"Use skip_installation_check=True to bypass this check."
|
|
)
|
|
# Don't raise an error, just warn - the parser might still work
|
|
|
|
def get_supported_extensions(self) -> List[str]:
|
|
"""Get list of supported file extensions"""
|
|
return list(
|
|
self.parser.OFFICE_FORMATS
|
|
| self.parser.IMAGE_FORMATS
|
|
| self.parser.TEXT_FORMATS
|
|
| {".pdf"}
|
|
)
|
|
|
|
def filter_supported_files(
|
|
self, file_paths: List[str], recursive: bool = True
|
|
) -> List[str]:
|
|
"""
|
|
Filter file paths to only include supported file types
|
|
|
|
Args:
|
|
file_paths: List of file paths or directories
|
|
recursive: Whether to search directories recursively
|
|
|
|
Returns:
|
|
List of supported file paths
|
|
"""
|
|
supported_extensions = set(self.get_supported_extensions())
|
|
supported_files = []
|
|
|
|
for path_str in file_paths:
|
|
path = Path(path_str)
|
|
|
|
if path.is_file():
|
|
if path.suffix.lower() in supported_extensions:
|
|
supported_files.append(str(path))
|
|
else:
|
|
self.logger.warning(f"Unsupported file type: {path}")
|
|
|
|
elif path.is_dir():
|
|
if recursive:
|
|
# Recursively find all files
|
|
for file_path in path.rglob("*"):
|
|
if (
|
|
file_path.is_file()
|
|
and file_path.suffix.lower() in supported_extensions
|
|
):
|
|
supported_files.append(str(file_path))
|
|
else:
|
|
# Only files in the directory (not subdirectories)
|
|
for file_path in path.glob("*"):
|
|
if (
|
|
file_path.is_file()
|
|
and file_path.suffix.lower() in supported_extensions
|
|
):
|
|
supported_files.append(str(file_path))
|
|
|
|
else:
|
|
self.logger.warning(f"Path does not exist: {path}")
|
|
|
|
return supported_files
|
|
|
|
def process_single_file(
|
|
self, file_path: str, output_dir: str, parse_method: str = "auto", **kwargs
|
|
) -> Tuple[bool, str, Optional[str]]:
|
|
"""
|
|
Process a single file
|
|
|
|
Args:
|
|
file_path: Path to the file to process
|
|
output_dir: Output directory
|
|
parse_method: Parsing method
|
|
**kwargs: Additional parser arguments
|
|
|
|
Returns:
|
|
Tuple of (success, file_path, error_message)
|
|
"""
|
|
try:
|
|
start_time = time.time()
|
|
|
|
# Create file-specific output directory
|
|
file_name = Path(file_path).stem
|
|
file_output_dir = Path(output_dir) / file_name
|
|
file_output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Parse the document
|
|
content_list = self.parser.parse_document(
|
|
file_path=file_path,
|
|
output_dir=str(file_output_dir),
|
|
method=parse_method,
|
|
**kwargs,
|
|
)
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
self.logger.info(
|
|
f"Successfully processed {file_path} "
|
|
f"({len(content_list)} content blocks, {processing_time:.2f}s)"
|
|
)
|
|
|
|
return True, file_path, None
|
|
|
|
except Exception as e:
|
|
error_msg = f"Failed to process {file_path}: {str(e)}"
|
|
self.logger.error(error_msg)
|
|
return False, file_path, error_msg
|
|
|
|
@staticmethod
|
|
def _manifest_path(output_dir: str) -> Path:
|
|
"""Return the incremental processing manifest path."""
|
|
return Path(output_dir) / ".raganything_batch_manifest.json"
|
|
|
|
def _load_incremental_manifest(
|
|
self, output_dir: str
|
|
) -> Dict[str, Dict[str, object]]:
|
|
"""Load incremental processing metadata."""
|
|
manifest_path = self._manifest_path(output_dir)
|
|
if not manifest_path.exists():
|
|
return {}
|
|
|
|
try:
|
|
data = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
self.logger.warning(
|
|
f"Could not read incremental manifest {manifest_path}: {exc}"
|
|
)
|
|
return {}
|
|
|
|
if not isinstance(data, dict):
|
|
self.logger.warning(
|
|
f"Ignoring invalid incremental manifest {manifest_path}: expected object"
|
|
)
|
|
return {}
|
|
|
|
files = data.get("files", {})
|
|
return files if isinstance(files, dict) else {}
|
|
|
|
def _save_incremental_manifest(
|
|
self, output_dir: str, manifest: Dict[str, Dict[str, object]]
|
|
) -> None:
|
|
"""Persist incremental processing metadata atomically."""
|
|
manifest_path = self._manifest_path(output_dir)
|
|
manifest_path.parent.mkdir(parents=True, exist_ok=True)
|
|
payload = {
|
|
"version": 1,
|
|
"updated_at": time.time(),
|
|
"files": manifest,
|
|
}
|
|
# Write to a uniquely named temp file in the same directory before the
|
|
# atomic replace. A fixed ".tmp" name would let two concurrent batches
|
|
# writing to the same output_dir clobber each other's half-written temp.
|
|
fd, temp_name = tempfile.mkstemp(
|
|
dir=str(manifest_path.parent),
|
|
prefix=".raganything_batch_manifest.",
|
|
suffix=".tmp",
|
|
)
|
|
temp_path = Path(temp_name)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as temp_file:
|
|
json.dump(payload, temp_file, indent=2, sort_keys=True)
|
|
temp_path.replace(manifest_path)
|
|
except Exception:
|
|
temp_path.unlink(missing_ok=True)
|
|
raise
|
|
|
|
@staticmethod
|
|
def _file_metadata(file_path: str) -> Dict[str, object]:
|
|
"""Return cheap stat metadata (resolved path, size, mtime) without hashing."""
|
|
path = Path(file_path)
|
|
stat = path.stat()
|
|
return {
|
|
"path": str(path.resolve()),
|
|
"size": stat.st_size,
|
|
"mtime_ns": stat.st_mtime_ns,
|
|
}
|
|
|
|
@staticmethod
|
|
def _compute_md5(file_path: str) -> str:
|
|
"""Stream a file and return its md5 hex digest."""
|
|
md5 = hashlib.md5()
|
|
with Path(file_path).open("rb") as file_obj:
|
|
for chunk in iter(lambda: file_obj.read(1024 * 1024), b""):
|
|
md5.update(chunk)
|
|
return md5.hexdigest()
|
|
|
|
@classmethod
|
|
def _file_signature(cls, file_path: str) -> Dict[str, object]:
|
|
"""Return size, mtime, and md5 metadata for change detection."""
|
|
metadata = cls._file_metadata(file_path)
|
|
metadata["md5"] = cls._compute_md5(file_path)
|
|
return metadata
|
|
|
|
def _filter_incremental_files(
|
|
self, file_paths: List[str], manifest: Dict[str, Dict[str, object]]
|
|
) -> Tuple[List[str], List[str], Dict[str, Dict[str, object]]]:
|
|
"""Split files into changed and unchanged groups using manifest metadata."""
|
|
files_to_process = []
|
|
skipped_files = []
|
|
signatures = {}
|
|
|
|
for file_path in file_paths:
|
|
try:
|
|
metadata = self._file_metadata(file_path)
|
|
except OSError as exc:
|
|
# File vanished or became unreadable between discovery and the
|
|
# incremental scan. Treat it as changed so the normal per-file
|
|
# error path handles it instead of aborting the whole batch.
|
|
self.logger.warning(
|
|
f"Could not read {file_path} for incremental scan: {exc}"
|
|
)
|
|
files_to_process.append(file_path)
|
|
continue
|
|
|
|
manifest_key = metadata["path"]
|
|
previous = manifest.get(manifest_key)
|
|
|
|
# Fast path: if size and mtime match a previous run, trust the
|
|
# stored signature and skip re-hashing the file entirely.
|
|
if (
|
|
isinstance(previous, dict)
|
|
and previous.get("size") == metadata["size"]
|
|
and previous.get("mtime_ns") == metadata["mtime_ns"]
|
|
and "md5" in previous
|
|
):
|
|
signatures[file_path] = previous
|
|
skipped_files.append(file_path)
|
|
continue
|
|
|
|
# Metadata differs (or the file is new/unrecorded): hash to build a
|
|
# full signature and compare content.
|
|
try:
|
|
metadata["md5"] = self._compute_md5(file_path)
|
|
except OSError as exc:
|
|
self.logger.warning(
|
|
f"Could not read {file_path} for incremental scan: {exc}"
|
|
)
|
|
files_to_process.append(file_path)
|
|
continue
|
|
|
|
signatures[file_path] = metadata
|
|
|
|
if previous == metadata:
|
|
skipped_files.append(file_path)
|
|
else:
|
|
files_to_process.append(file_path)
|
|
|
|
return files_to_process, skipped_files, signatures
|
|
|
|
def process_batch(
|
|
self,
|
|
file_paths: List[str],
|
|
output_dir: str,
|
|
parse_method: str = "auto",
|
|
recursive: bool = True,
|
|
dry_run: bool = False,
|
|
incremental: bool = False,
|
|
**kwargs,
|
|
) -> BatchProcessingResult:
|
|
"""
|
|
Process multiple files in parallel
|
|
|
|
Args:
|
|
file_paths: List of file paths or directories to process
|
|
output_dir: Base output directory
|
|
parse_method: Parsing method for all files
|
|
recursive: Whether to search directories recursively
|
|
dry_run: When True, only list files without processing them
|
|
incremental: When True, skip files whose size, mtime, and md5 match
|
|
the previous successful batch run in output_dir
|
|
**kwargs: Additional parser arguments
|
|
|
|
Returns:
|
|
BatchProcessingResult with processing statistics
|
|
"""
|
|
start_time = time.time()
|
|
|
|
# Filter to supported files
|
|
supported_files = self.filter_supported_files(file_paths, recursive)
|
|
|
|
if not supported_files:
|
|
self.logger.warning("No supported files found to process")
|
|
return BatchProcessingResult(
|
|
successful_files=[],
|
|
failed_files=[],
|
|
total_files=0,
|
|
processing_time=0.0,
|
|
errors={},
|
|
output_dir=output_dir,
|
|
dry_run=dry_run,
|
|
)
|
|
|
|
self.logger.info(f"Found {len(supported_files)} files to process")
|
|
|
|
# Create output directory before reading/writing the incremental manifest
|
|
output_path = Path(output_dir)
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
skipped_files: List[str] = []
|
|
manifest: Dict[str, Dict[str, object]] = {}
|
|
signatures: Dict[str, Dict[str, object]] = {}
|
|
files_to_process = supported_files
|
|
if incremental:
|
|
manifest = self._load_incremental_manifest(output_dir)
|
|
files_to_process, skipped_files, signatures = (
|
|
self._filter_incremental_files(supported_files, manifest)
|
|
)
|
|
self.logger.info(
|
|
f"Incremental scan: {len(files_to_process)} changed files, "
|
|
f"{len(skipped_files)} unchanged files"
|
|
)
|
|
|
|
if dry_run:
|
|
self.logger.info(
|
|
f"Dry run enabled. {len(files_to_process)} files would be processed."
|
|
)
|
|
return BatchProcessingResult(
|
|
successful_files=files_to_process,
|
|
failed_files=[],
|
|
total_files=len(supported_files),
|
|
processing_time=0.0,
|
|
errors={},
|
|
output_dir=output_dir,
|
|
dry_run=True,
|
|
skipped_files=skipped_files,
|
|
)
|
|
|
|
# Process files in parallel
|
|
successful_files = []
|
|
failed_files = []
|
|
errors = {}
|
|
|
|
# Create progress bar if requested
|
|
pbar = None
|
|
if self.show_progress:
|
|
pbar = tqdm(
|
|
total=len(files_to_process),
|
|
desc=f"Processing files ({self.parser_type})",
|
|
unit="file",
|
|
)
|
|
|
|
future_to_file = {}
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
|
|
# Submit all tasks
|
|
future_to_file = {
|
|
executor.submit(
|
|
self.process_single_file,
|
|
file_path,
|
|
output_dir,
|
|
parse_method,
|
|
**kwargs,
|
|
): file_path
|
|
for file_path in files_to_process
|
|
}
|
|
|
|
# Process completed tasks
|
|
for future in as_completed(
|
|
future_to_file, timeout=self.timeout_per_file
|
|
):
|
|
success, file_path, error_msg = future.result()
|
|
|
|
if success:
|
|
successful_files.append(file_path)
|
|
else:
|
|
failed_files.append(file_path)
|
|
errors[file_path] = error_msg
|
|
|
|
if pbar:
|
|
pbar.update(1)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Batch processing failed: {str(e)}")
|
|
# Mark remaining files as failed
|
|
for future in future_to_file:
|
|
if not future.done():
|
|
file_path = future_to_file[future]
|
|
failed_files.append(file_path)
|
|
errors[file_path] = f"Processing interrupted: {str(e)}"
|
|
if pbar:
|
|
pbar.update(1)
|
|
|
|
finally:
|
|
if pbar:
|
|
pbar.close()
|
|
|
|
if incremental and successful_files:
|
|
for file_path in successful_files:
|
|
signature = signatures.get(file_path)
|
|
if signature:
|
|
manifest[signature["path"]] = signature
|
|
self._save_incremental_manifest(output_dir, manifest)
|
|
|
|
processing_time = time.time() - start_time
|
|
|
|
# Create result
|
|
result = BatchProcessingResult(
|
|
successful_files=successful_files,
|
|
failed_files=failed_files,
|
|
total_files=len(supported_files),
|
|
processing_time=processing_time,
|
|
errors=errors,
|
|
output_dir=output_dir,
|
|
dry_run=False,
|
|
skipped_files=skipped_files,
|
|
)
|
|
|
|
# Log summary
|
|
self.logger.info(result.summary())
|
|
|
|
return result
|
|
|
|
async def process_batch_async(
|
|
self,
|
|
file_paths: List[str],
|
|
output_dir: str,
|
|
parse_method: str = "auto",
|
|
recursive: bool = True,
|
|
dry_run: bool = False,
|
|
incremental: bool = False,
|
|
**kwargs,
|
|
) -> BatchProcessingResult:
|
|
"""
|
|
Async version of batch processing
|
|
|
|
Args:
|
|
file_paths: List of file paths or directories to process
|
|
output_dir: Base output directory
|
|
parse_method: Parsing method for all files
|
|
recursive: Whether to search directories recursively
|
|
dry_run: When True, only list files without processing them
|
|
incremental: When True, skip unchanged files based on the batch manifest
|
|
**kwargs: Additional parser arguments
|
|
|
|
Returns:
|
|
BatchProcessingResult with processing statistics
|
|
"""
|
|
# Run the sync version in a thread pool
|
|
loop = asyncio.get_event_loop()
|
|
process_func = partial(
|
|
self.process_batch,
|
|
file_paths=file_paths,
|
|
output_dir=output_dir,
|
|
parse_method=parse_method,
|
|
recursive=recursive,
|
|
dry_run=dry_run,
|
|
incremental=incremental,
|
|
**kwargs,
|
|
)
|
|
return await loop.run_in_executor(None, process_func)
|
|
|
|
|
|
def main():
|
|
"""Command-line interface for batch parsing"""
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Batch document parsing")
|
|
parser.add_argument("paths", nargs="+", help="File paths or directories to process")
|
|
parser.add_argument("--output", "-o", required=True, help="Output directory")
|
|
parser.add_argument(
|
|
"--parser",
|
|
default="mineru",
|
|
help=(
|
|
"Parser to use. Built-ins: mineru, docling, paddleocr. "
|
|
"When using RAGAnything as a library, any custom parsers that you "
|
|
"have registered via register_parser() in the current process "
|
|
"are also accepted. The standalone CLI itself does not perform "
|
|
"plugin discovery."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--method",
|
|
choices=["auto", "txt", "ocr"],
|
|
default="auto",
|
|
help="Parsing method",
|
|
)
|
|
parser.add_argument(
|
|
"--workers", type=int, default=4, help="Number of parallel workers"
|
|
)
|
|
parser.add_argument(
|
|
"--no-progress", action="store_true", help="Disable progress bar"
|
|
)
|
|
parser.add_argument(
|
|
"--recursive",
|
|
action="store_true",
|
|
default=True,
|
|
help="Search directories recursively",
|
|
)
|
|
parser.add_argument(
|
|
"--timeout", type=int, default=300, help="Timeout per file (seconds)"
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
help="List files that would be processed without running parsers",
|
|
)
|
|
parser.add_argument(
|
|
"--incremental",
|
|
action="store_true",
|
|
help="Skip files unchanged since the last successful batch run",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
)
|
|
|
|
try:
|
|
# Create batch parser
|
|
batch_parser = BatchParser(
|
|
parser_type=args.parser,
|
|
max_workers=args.workers,
|
|
show_progress=not args.no_progress,
|
|
timeout_per_file=args.timeout,
|
|
)
|
|
|
|
# Process files
|
|
result = batch_parser.process_batch(
|
|
file_paths=args.paths,
|
|
output_dir=args.output,
|
|
parse_method=args.method,
|
|
recursive=args.recursive,
|
|
dry_run=args.dry_run,
|
|
incremental=args.incremental,
|
|
)
|
|
|
|
# Print summary
|
|
print("\n" + result.summary())
|
|
|
|
if args.dry_run:
|
|
if result.successful_files:
|
|
print("\nDry run: files that would be processed:")
|
|
for file_path in result.successful_files:
|
|
print(f" - {file_path}")
|
|
else:
|
|
print("\nDry run: no supported files found.")
|
|
if result.skipped_files:
|
|
print("\nDry run: unchanged files that would be skipped:")
|
|
for file_path in result.skipped_files:
|
|
print(f" - {file_path}")
|
|
|
|
# Exit with error code if any files failed
|
|
if result.failed_files:
|
|
return 1
|
|
|
|
return 0
|
|
|
|
except Exception as e:
|
|
print(f"Error: {str(e)}")
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|