# 文件名: processor.py import os import sys import contextlib from io import StringIO import json from pathlib import Path from datetime import datetime import re import time import traceback import streamlit as st import ebooklib import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM from book_maker.loader import BOOK_LOADER_DICT from book_maker.translator import MODEL_DICT from book_maker.utils import LANGUAGES, prompt_config_to_kwargs # --- 本地模型函数 --- @st.cache_resource def load_helsinki_model(model_name): try: tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForSeq2SeqLM.from_pretrained(model_name) return tokenizer, model except Exception as e: raise RuntimeError(f"Helsinki Model loading failed: {e}") def helsinki_translate_paragraphs(paragraphs, language, direction_str, status_container): if not direction_str or "en-zh" in direction_str: model_name = "Helsinki-NLP/opus-mt-en-zh" elif "zh-en" in direction_str: model_name = "Helsinki-NLP/opus-mt-zh-en" else: raise ValueError(f"未知的 Helsinki-NLP 翻译方向: {direction_str}") tokenizer, model = load_helsinki_model(model_name) translated_paragraphs = [] pbar_container = status_container.container() pbar = pbar_container.progress(0, text="Helsinki模型翻译进度: 0%") for i, para in enumerate(paragraphs): if para.strip(): inputs = tokenizer(para, return_tensors="pt", padding=True, truncation=True, max_length=512) translated_tokens = model.generate(**inputs) translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0] translated_paragraphs.append(translated_text) else: translated_paragraphs.append("") progress = (i + 1) / len(paragraphs) if pbar: pbar.progress(progress, text=f"Helsinki模型翻译进度: {i+1}/{len(paragraphs)}") return translated_paragraphs # --- 通用工具 --- @contextlib.contextmanager def st_redirect(streamlit_element): original_stdout = sys.stdout output_catcher = StringIO() sys.stdout = output_catcher try: yield finally: sys.stdout = original_stdout # --- 【核心】独立的任务执行器 --- def run_translation_task(task_id, input_file, engine_id, language, api_key, kwargs_options): DB_FILE = Path("tasks_db.json") def update_task_status(status, result_file=None, error_message=None): tasks = {} if DB_FILE.exists(): lock_file = DB_FILE.with_suffix(".lock") while lock_file.exists(): time.sleep(0.1) try: lock_file.touch() with open(DB_FILE, "r", encoding="utf-8") as f: tasks = json.load(f) finally: if lock_file.exists(): lock_file.unlink() if task_id in tasks: tasks[task_id]["status"] = status tasks[task_id]["updated_at"] = datetime.now().isoformat() if result_file: tasks[task_id]["result_file"] = str(result_file) if error_message: tasks[task_id]["error"] = error_message lock_file = DB_FILE.with_suffix(".lock") while lock_file.exists(): time.sleep(0.1) try: lock_file.touch() with open(DB_FILE, "w", encoding="utf-8") as f: json.dump(tasks, f, indent=2, ensure_ascii=False) finally: if lock_file.exists(): lock_file.unlink() try: update_task_status("🏃‍♂️ 运行中...") class MockStreamlitElement: def empty(self): return self def container(self): return self def progress(self, *args, **kwargs): pass def text(self, *args, **kwargs): pass def info(self, *args, **kwargs): pass mock_status_container = MockStreamlitElement() output_file = translate_book_processing(input_file, engine_id, api_key, language, mock_status_container, **kwargs_options) update_task_status("✅ 已完成", result_file=output_file) except Exception as e: error_details = traceback.format_exc() update_task_status("❌ 失败", error_message=f"{str(e)}\n\nTraceback:\n{error_details}") # --- 【核心】原始翻译处理函数 --- def translate_book_processing(input_file_path, engine_id, api_key, language, status_container, **kwargs): book_type = input_file_path.split('.')[-1] # 【修复】从带有时间戳的输入文件名中，恢复出原始文件名 # 例如：从 "100849_【初祥】发烧.txt" 变为 "【初祥】发烧.txt" original_name_with_ext = "_".join(Path(input_file_path).name.split('_')[1:]) name_only, ext_only = os.path.splitext(original_name_with_ext) output_dir = Path("outputs") output_dir.mkdir(exist_ok=True) # 【修复】使用恢复后的原始文件名来创建双语文件名 output_file_path = str(output_dir / f"{name_only}_bilingual{ext_only}") if engine_id == "local_helsinki": single_translate = kwargs.get("single_translate", False) book_loader_class = BOOK_LOADER_DICT.get(book_type) if not book_loader_class: raise ValueError(f"不支持的文件格式: {book_type}") temp_loader_options = {"model": lambda k, l, **kw: None, "key": "", "language": language, "resume": kwargs.get("resume", False)} temp_loader_options[f"{book_type}_name"] = input_file_path book = book_loader_class(**temp_loader_options) if book_type == "epub": new_book = book._make_new_book(book.origin_book) all_items = list(book.origin_book.get_items_of_type(ebooklib.ITEM_DOCUMENT)) for item in book.origin_book.get_items(): if item.get_type() != ebooklib.ITEM_DOCUMENT: new_book.add_item(item) for item in all_items: soup = book.bs(item.content, "html.parser") tags_to_translate = kwargs.get("translate_tags", "p").split(",") p_list = soup.findAll(tags_to_translate) original_texts = [p.get_text() for p in p_list] translated_texts = helsinki_translate_paragraphs(original_texts, language, kwargs.get("direction"), status_container) book.helper = book_loader_class.helper_class(translate_model=None) for p, trans in zip(p_list, translated_texts): book.helper.insert_trans(p, trans, "", single_translate) item.content = soup.encode() new_book.add_item(item) ebooklib.epub.write_epub(output_file_path, new_book, {}) else: with open(input_file_path, 'r', encoding='utf-8') as f: paragraphs = f.read().split('\n\n') translated_paragraphs = helsinki_translate_paragraphs(paragraphs, language, kwargs.get("direction"), status_container) with open(output_file_path, 'w', encoding='utf-8') as f: for orig, trans in zip(paragraphs, translated_paragraphs): if not single_translate: f.write(orig + '\n\n') f.write(trans + '\n\n') return output_file_path translate_model_class = MODEL_DICT.get(engine_id) if not translate_model_class: raise ValueError(f"不支持的翻译引擎: {engine_id}") loader_arg_name = f"{book_type}_name" language_value = LANGUAGES.get(language, language) loader_options = { loader_arg_name: input_file_path, "model": translate_model_class, "key": api_key or "None", "language": language_value, "resume": kwargs.get("resume", False), "is_test": kwargs.get("is_test", False), "test_num": kwargs.get("test_num", 10), "prompt_config": kwargs.get("prompt_config"), "single_translate": kwargs.get("single_translate", False), "context_flag": kwargs.get("context_flag", False), "temperature": kwargs.get("temperature", 1.0), "source_lang": kwargs.get("source_lang", "auto"), "model_api_base": kwargs.get("api_base"), } proxy = kwargs.get('proxy') if proxy: os.environ["http_proxy"] = os.environ["https_proxy"] = proxy book_translator = BOOK_LOADER_DICT.get(book_type)(**loader_options) with st_redirect(status_container): book_translator.make_bilingual_book() if os.path.exists(output_file_path): return output_file_path # Fallback check for original library's naming convention # e.g., for input 'uploads/123_test.txt', it might create 'uploads/123_test_bilingual.txt' temp_output_path = f"{os.path.splitext(input_file_path)[0]}_bilingual{ext_only}" if os.path.exists(temp_output_path): os.rename(temp_output_path, output_file_path) return output_file_path raise FileNotFoundError(f"翻译完成，但未在预期的输出路径找到文件: {output_file_path}")