import argparse import hashlib import json import subprocess 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.glossary import prepare_glossary 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.") llm_client = None 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"]) if structure.source_sha256 != hashlib.sha256(input_path.read_bytes()).hexdigest(): logger.info("Source changed or old preprocessing cache; re-cleaning") structure = BookStructure.load(EpubCleaner(input_path, work["root"]).clean()) 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), timeout=llm_conf.get("timeout", 180), max_output_tokens=llm_conf.get("max_output_tokens", 16000), 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, cache_path=work["root"] / "profile.json") 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) glossary = await prepare_glossary( manifest_manager.entries, llm_client, work["root"] / "glossary.json", override=getattr(args, 'glossary', None) or trans_conf.get("glossary_path"), auto=trans_conf.get("auto_glossary", True), chunk_chars=trans_conf.get("glossary_scan_chars", 24000)) translator = Translator(llm_client, chunk_size=target_chunk_size, max_concurrent=concurrent_reqs, glossary=glossary, validation_retries=trans_conf.get("validation_retries", 2), on_chunk_complete=manifest_manager.save) await translator.translate(manifest_manager.entries, profile) manifest_manager.save() if any(not e.translated_text or e.validation_errors for e in manifest_manager.entries): raise ValueError("Translation incomplete; inspect manifest validation_errors and resume") 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.") if any(e.translated_text and not e.translated_html for e in manifest_manager.entries): raise ValueError("Format validation failed; broken entries were not backfilled") # 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, allow_partial=args.skip_translation) # 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) from src.qa.epub_audit import audit_epub, compare_audits source_audit, output_audit = audit_epub(input_path), audit_epub(created_epub) qa_report = {'source': source_audit, 'output': output_audit, 'new_issues': compare_audits(source_audit, output_audit), 'visual_review': 'pending'} (work["root"] / "structural-qa.json").write_text(json.dumps(qa_report, ensure_ascii=False, indent=2)) if qa_report['new_issues']: raise ValueError("EPUB structural QA failed; inspect structural-qa.json") if getattr(args, 'visual_qa', False): command = [sys.executable, str(Path(__file__).parent / 'scripts/visual_qa.py'), str(input_path), str(created_epub), '--manifest', str(work['manifest']), '--out', str(work['root'] / 'visual-qa'), '--render'] if args.skip_translation: command.append('--allow-partial') subprocess.run(command, check=True) logger.info(f"Pipeline completed! Output: {created_epub}") 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) finally: if llm_client: await llm_client.close() 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("--visual-qa", action="store_true", help="Render Calibre QA page images for visual inspection") parser.add_argument("--glossary", help="Per-book JSON terminology file") 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()