87 lines
2.9 KiB
Python
87 lines
2.9 KiB
Python
import argparse
|
|
import sys
|
|
import os
|
|
import asyncio
|
|
from pathlib import Path
|
|
from dotenv import load_dotenv
|
|
|
|
# Add project root to sys.path
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from src.common.utils import setup_logger, ensure_directory
|
|
from src.common.paths import get_work_dirs
|
|
from src.common.config import load_global_config
|
|
from src.translation.manifest_manager import ManifestManager
|
|
from src.translation.llm_client import LLMClient
|
|
from src.assembly.restoration_engine import RestorationEngine
|
|
|
|
logger = setup_logger("pipeline_restore")
|
|
|
|
async def run_restore(args):
|
|
input_path = Path(args.input_epub)
|
|
if not input_path.exists():
|
|
logger.error(f"Input file not found: {input_path}")
|
|
sys.exit(1)
|
|
|
|
work = get_work_dirs(input_path)
|
|
manifest_path = work["manifest"]
|
|
|
|
if not manifest_path.exists():
|
|
logger.error("Missing manifest. Run Step 1 & 2.")
|
|
sys.exit(1)
|
|
|
|
# Load Config
|
|
load_dotenv()
|
|
config = load_global_config()
|
|
llm_conf = config.get("llm", {})
|
|
api_key = llm_conf.get("api_key") or os.getenv("OPENAI_API_KEY")
|
|
|
|
llm_client = None
|
|
if api_key:
|
|
logger.info("Initializing LLM Client for repairs...")
|
|
llm_client = LLMClient(
|
|
api_key=api_key,
|
|
base_url=llm_conf.get("base_url"),
|
|
model=llm_conf.get("model", "gpt-3.5-turbo"),
|
|
requests_per_minute=llm_conf.get("requests_per_minute", 60),
|
|
concurrent_requests=llm_conf.get("concurrent_requests", 5),
|
|
chunk_dir=work["chunks"] # Reuse chunks dir for logging repairs
|
|
)
|
|
else:
|
|
logger.warning("No API Key. LLM Repair disabled.")
|
|
|
|
try:
|
|
# Load Manifest
|
|
manager = ManifestManager(manifest_path)
|
|
manager.load()
|
|
logger.info(f"Loaded {len(manager.entries)} entries.")
|
|
|
|
# Restore Phase
|
|
engine = RestorationEngine(llm_client)
|
|
logger.info("Starting format restoration (Spacing + Tags + Repair)...")
|
|
if args.force:
|
|
logger.info("Force mode enabled: Re-processing all entries.")
|
|
|
|
success_count = await engine.restore_entries(manager.entries, force=args.force)
|
|
|
|
# Save Result
|
|
manager.save()
|
|
logger.info(f"Restoration complete. {success_count}/{len(manager.entries)} fully validated.")
|
|
|
|
except Exception as e:
|
|
logger.critical(f"Restoration failed: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
sys.exit(1)
|
|
finally:
|
|
if llm_client:
|
|
await llm_client.close()
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Step 3: Restore Format (HTML Generation)")
|
|
parser.add_argument("input_epub", help="Path to original EPUB")
|
|
parser.add_argument("--force", action="store_true", help="Force re-restoration")
|
|
|
|
args = parser.parse_args()
|
|
asyncio.run(run_restore(args))
|