import argparse import sys import os import asyncio from pathlib import Path from dotenv import load_dotenv from src.common.config import load_global_config from src.preprocessing.epub_cleaner import EpubCleaner from src.preprocessing.profiler import BookProfiler from src.preprocessing.text_extractor import FineGrainedExtractor from src.translation.manifest_manager import ManifestManager from src.translation.translator_engine import Translator from src.translation.llm_client import LLMClient from src.assembly.restoration_engine import RestorationEngine from src.assembly.backfiller import BackfillEngine from src.assembly.builder import BilingualBuilder from src.common.data_model import BookStructure from src.common.utils import setup_logger, ensure_directory from src.common.exceptions import EpubTranslatorError logger = setup_logger("main") # Unified work directory structure # .work/ # ├── {book_name}/ # │ ├── book_structure.json # │ ├── manifest.json # │ ├── assets/ # │ └── chunks/ from src.common.paths import get_work_dirs async def run_pipeline(args): input_path = Path(args.input_epub) output_dir = Path(args.output_dir) # Get work directories for this book work = get_work_dirs(input_path) ensure_directory(work["root"]) ensure_directory(output_dir) # Load Config config = load_global_config() llm_conf = config.get("llm", {}) trans_conf = config.get("translation", {}) api_key = llm_conf.get("api_key") # Base URL and Model come from config if not overridden base_url = llm_conf.get("base_url") # CLI model arg overrides config model, which overrides default model = args.model if args.model else llm_conf.get("model", "gpt-3.5-turbo") if not api_key: logger.warning("OPENAI_API_KEY not found in env or config. LLM features may fail.") try: # 1. Preprocessing - Reuse book_structure.json if exists if work["structure"].exists() and not args.force_clean: logger.info(f"Reusing existing book_structure: {work['structure']}") structure = BookStructure.load(work["structure"]) else: logger.info("Cleaning EPUB and generating book_structure...") cleaner = EpubCleaner(input_path, work["root"]) book_structure_json = cleaner.clean() structure = BookStructure.load(book_structure_json) # 2. Extraction extractor = FineGrainedExtractor() manifest_entries = extractor.extract(structure) # 3. Manifest Management manifest_manager = ManifestManager(work["manifest"]) manifest_manager.load() # Load existing if any manifest_manager.add_entries(manifest_entries) manifest_manager.save() # 4. Translation # Initialize LLM Client (Centralized) llm_client = None if api_key: llm_client = LLMClient( api_key=api_key, base_url=base_url, model=model, requests_per_minute=llm_conf.get("requests_per_minute", 60), concurrent_requests=llm_conf.get("concurrent_requests", 5), chunk_dir=work["chunks"] ) # 4. Translation if not args.skip_translation: if not llm_client: logger.error("Cannot translate without API Key. Use --skip-translation to test pipeline.") sys.exit(1) # Profiling profiler = BookProfiler(llm_client) profile = await profiler.analyze(manifest_manager.entries) logger.info(f"Book Profile: {profile}") # Translation target_chunk_size = trans_conf.get("chunk_size", 5000) concurrent_reqs = llm_conf.get("concurrent_requests", 5) translator = Translator(llm_client, chunk_size=target_chunk_size, max_concurrent=concurrent_reqs) await translator.translate(manifest_manager.entries, profile) manifest_manager.save() else: logger.info("Skipping translation step.") # 5. Restoration (Format + Spacing + Repair) logger.info("Restoring format & applying spacing...") # RestorationEngine handles Pangu spacing, FormatRestorer, and LLM Repair restorer = RestorationEngine(llm_client) # Default to skipping if already done, unless forced force_restore = getattr(args, 'force_restore', False) restore_success = await restorer.restore_entries(manifest_manager.entries, force=force_restore) manifest_manager.save() # Save translated_html logger.info(f"Restoration complete. {restore_success} entries validated.") # 6. Backfill (Pure Injection) logger.info("Injecting content into EPUB structure...") backfiller = BackfillEngine() # Pure injection, no dependencies updated_structure = await backfiller.backfill(structure, manifest_manager.entries, mode=args.mode) # 7. Assembly - Pass original EPUB for TOC preservation builder = BilingualBuilder(work["root"], original_epub_path=input_path) if args.mode == "bilingual": output_filename = f"bilingual_{input_path.name}" else: output_filename = f"translated_{input_path.name}" output_path = output_dir / output_filename created_epub = builder.build(updated_structure, output_path) logger.info(f"Pipeline completed! Output: {created_epub}") if llm_client: await llm_client.close() except EpubTranslatorError as e: logger.error(f"An error occurred: {e}") sys.exit(1) except Exception as e: logger.critical(f"Unexpected error: {e}") import traceback traceback.print_exc() sys.exit(1) def main(): parser = argparse.ArgumentParser(description="EPUB Bilingual Translator") parser.add_argument("input_epub", help="Path to the input EPUB file") parser.add_argument("--output-dir", default="output", help="Directory for output files") parser.add_argument("--model", default=None, help="LLM Model to use (overrides config)") parser.add_argument("--bilingual", action="store_true", help="Output bilingual version (default is target language only)") parser.add_argument("--skip-translation", action="store_true", help="Skip LLM translation (for testing)") parser.add_argument("--force-restore", action="store_true", help="Force re-run format restoration/repair even if translated_html exists") parser.add_argument("--force-clean", action="store_true", help="Force re-clean EPUB even if book_structure exists") args = parser.parse_args() # Map boolean flag to mode string args.mode = "bilingual" if args.bilingual else "target_only" asyncio.run(run_pipeline(args)) if __name__ == "__main__": main()