import os import gc import gradio as gr import torch import random import logging from huggingface_hub import login, HfApi, snapshot_download import spacy import subprocess import pkg_resources import sys login(token=os.environ.get("LA_NAME")) os.environ["LASER"] = "laser" def check_and_install(package, required_version): try: dist = pkg_resources.get_distribution(package) installed_version = dist.version if installed_version != required_version: print(f"[{package}] already installed {installed_version}. Required version {required_version}，re-install...") subprocess.check_call([sys.executable, "-m", "pip", "install", f"{package}=={required_version}", "--force-reinstall"]) else: print(f"[{package}] required version {required_version} finished") except pkg_resources.DistributionNotFound: print(f"[{package}] not found, install: {required_version}...") subprocess.check_call([sys.executable, "-m", "pip", "install", f"{package}=={required_version}"]) packages = { "pip": "24.0", "fairseq": "0.12.2", "torch": "2.6.0", "transformers": "4.51.3" } for package, version in packages.items(): check_and_install(package, version) from transformers import AutoTokenizer, AutoModelForCausalLM from vecalign.plan2align import translate_text, external_find_best_translation from trl import AutoModelForCausalLMWithValueHead models = ["en_core_web_sm", "ru_core_news_sm", "de_core_news_sm", "ja_core_news_sm", "ko_core_news_sm", "es_core_news_sm"] for model in models: try: spacy.load(model) except OSError: from spacy.cli import download download(model) try: spacy.load("zh_core_web_sm") except OSError: from spacy.cli import download download("zh_core_web_sm") subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy==1.24.0", "--force-reinstall"]) # ---------- translation function ---------- # Initialize device device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Load models once print("Loading models...") model_id = "google/gemma-2-9b-it" # "meta-llama/Meta-Llama-3.1-8B-Instruct" tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained( model_id, device_map="auto", torch_dtype=torch.float16 ) import spacy lang_map = { "English": ("en", "en_core_web_sm"), "Russian": ("ru", "ru_core_news_sm"), "German": ("de", "de_core_news_sm"), "Japanese": ("ja", "ja_core_news_sm"), "Korean": ("ko", "ko_core_news_sm"), "Spanish": ("es", "es_core_news_sm"), "Simplified Chinese": ("zh", "zh_core_web_sm"), "Traditional Chinese": ("zh", "zh_core_web_sm") } def get_lang_and_nlp(language): if language not in lang_map: raise ValueError(f"Unsupported language: {language}") lang_code, model_name = lang_map[language] return lang_code, spacy.load(model_name) def segment_sentences_by_punctuation(text, src_nlp): segmented_sentences = [] paragraphs = text.split('\n') for paragraph in paragraphs: if paragraph.strip(): doc = src_nlp(paragraph) for sent in doc.sents: segmented_sentences.append(sent.text.strip()) return segmented_sentences def generate_translation(system_prompt, prompt): full_prompt = f"System: {system_prompt}\nUser: {prompt}\nAssistant:" inputs = tokenizer(full_prompt, return_tensors="pt").to(device) outputs = model.generate( **inputs, max_new_tokens=2048, temperature=0.7, top_p=0.9, do_sample=True ) translation = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True) return translation def check_token_length(text, max_tokens=1024): return len(text) <= max_tokens import uuid def get_user_session(state=None): if state is None: state = {} if not isinstance(state, dict): state = {} if not state.get("session_id"): state["session_id"] = uuid.uuid4().hex return state["session_id"] # ---------- Translation Function ---------- def mpc_initial_translate(source_sentence, src_language, tgt_language): system_prompts = [ "You are a meticulous translator. Provide a literal, word-for-word translation that preserves the structure and meaning of each individual word.", "You are a professional translator. Deliver a clear, formal, and precise translation that faithfully conveys the original meaning.", "You are a creative and expressive translator. Render the text in a vivid and imaginative way, as if narrating a captivating story." ] translations = [] for prompt_style in system_prompts: prompt = f"### Translate this from {src_language} to {tgt_language} and only output the result." prompt += f"\n### {src_language}:\n {source_sentence}" prompt += f"\n### {tgt_language}:\n" translation = generate_translation(prompt_style, prompt) translations.append(translation) print("mpc_initial_translate") print(translations) return translations def mpc_improved_translate(source_sentence, current_translation, src_language, tgt_language): system_prompts = [ "You are a meticulous translator. Please improve the following translation by ensuring it is a literal and structurally precise version.", "You are a professional translator. Please refine the provided translation to be clear, formal, and accurate.", "You are a creative translator. Please enhance the translation so that it is vivid, natural, and engaging." ] translations = [] for prompt_style in system_prompts: prompt = (f"Source ({src_language}): {source_sentence}\n" f"Current Translation ({tgt_language}): {current_translation}\n" f"Please provide an improved translation into {tgt_language} and only output the result:") translation = generate_translation(prompt_style, prompt) translations.append(translation) print("mpc_improved_translate") print(translations) return translations def basic_translate(source_sentence, src_language, tgt_language): system_prompts = ["You are a helpful translator and only output the result."] translations = [] for prompt_style in system_prompts: prompt = f"### Translate this from {src_language} to {tgt_language}." prompt += f"\n### {src_language}:\n {source_sentence}" prompt += f"\n### {tgt_language}:\n" translation = generate_translation(prompt_style, prompt) translations.append(translation) return translations def summary_translate(src_text, temp_tgt_text, tgt_language, session_id): if len(temp_tgt_text.strip()) == 0: return "", 0 system_prompts = ["You are a helpful rephraser. You only output the rephrased result."] translations = [] for prompt_style in system_prompts: prompt = f"### Rephrase the following in {tgt_language}." prompt += f"\n### Input:\n {temp_tgt_text}" prompt += f"\n### Rephrased:\n" translation = generate_translation(prompt_style, prompt) translations.append(translation) try: _, score = evaluate_candidates(src_text, translations, tgt_language, session_id) except: score = 0 return translations[0], score def plan2align_translate_text(text, session_id, model, tokenizer, device, src_language, task_language, max_iterations_value, threshold_value, good_ref_contexts_num_value, reward_model_type): result = translate_text( text = text, model = model, tokenizer = tokenizer, device = device, src_language=src_language, task_language=task_language, max_iterations_value=max_iterations_value, threshold_value=threshold_value, good_ref_contexts_num_value=good_ref_contexts_num_value, reward_model_type=reward_model_type, session_id=session_id ) try: _, score = evaluate_candidates(text, [result], task_language, session_id) except: score = 0 return result, score def evaluate_candidates(source, candidates, language, session_id): evals = [(source, candidates)] best_translations = external_find_best_translation(evals, language, session_id) best_candidate, best_score = best_translations[0] return best_candidate, best_score def original_translation(text, src_language, target_language, session_id): cand_list = basic_translate(text, src_language, target_language) best, score = evaluate_candidates(text, cand_list, target_language, session_id) if cand_list: return best, score return "", 0 def best_of_n_translation(text, src_language, target_language, n, session_id): if not check_token_length(text, 4096): return "Warning: Input text too long.", 0 candidates = [] for i in range(n): cand_list = basic_translate(text, src_language, target_language) if cand_list: candidates.append(cand_list[0]) try: best, score = evaluate_candidates(text, candidates, target_language, session_id) print("best_of_n evaluate_candidates results:") print(best, score) except: print("evaluate_candidates fail") return "Warning: Input text too long.", 0 return best, score def mpc_translation(text, src_language, target_language, iterations, session_id): if not check_token_length(text, 4096): return "Warning: Input text too long.", 0 current_trans = "" best_score = None for i in range(iterations): if i == 0: cand_list = mpc_initial_translate(text, src_language, target_language) else: cand_list = mpc_improved_translate(text, current_trans, src_language, target_language) try: best, score = evaluate_candidates(text, cand_list, target_language, session_id) print("mpc evaluate_candidates results:") print(best, score) current_trans = best best_score = score except: print("evaluate_candidates fail") current_trans = cand_list[0] best_score = 0 return current_trans, best_score # ---------- Gradio function ---------- def process_text(text, src_language, target_language, max_iterations_value, threshold_value, good_ref_contexts_num_value, translation_methods=None, chunk_size=-1, state=None): translation_methods = translation_methods or ["Original", "Plan2Align"] session_id = get_user_session(state) """ 傳入中文文本與目標語言，依序產生四種翻譯結果： 1. 原始翻譯 2. Plan2Align 翻譯 3. Best-of-N 翻譯 4. MPC 翻譯 """ orig_output = "" plan2align_output = "" best_of_n_output = "" mpc_output = "" src_lang, src_nlp = get_lang_and_nlp(src_language) source_sentence = text.replace("\n", " ") source_segments = segment_sentences_by_punctuation(source_sentence, src_nlp) if chunk_size == -1: if "Original" in translation_methods: orig, best_score = original_translation(text, src_language, target_language, session_id) orig_output = f"{orig}\n\nScore: {best_score:.2f}" if "Plan2Align" in translation_methods: plan2align_trans, best_score = plan2align_translate_text( text, session_id, model, tokenizer, device, src_language, target_language, max_iterations_value, threshold_value, good_ref_contexts_num_value, "metricx" ) plan2align_output = f"{plan2align_trans}\n\nScore: {best_score:.2f}" if "Best-of-N" in translation_methods: best_candidate, best_score = best_of_n_translation(text, src_language, target_language, max_iterations_value, session_id) best_of_n_output = f"{best_candidate}\n\nScore: {best_score:.2f}" if "MPC" in translation_methods: mpc_candidate, mpc_score = mpc_translation(text, src_language, target_language, max_iterations_value, session_id) mpc_output = f"{mpc_candidate}\n\nScore: {mpc_score:.2f}" else: chunks = [' '.join(source_segments[i:i+chunk_size]) for i in range(0, len(source_segments), chunk_size)] org_translated_chunks = [] p2a_translated_chunks = [] bfn_translated_chunks = [] mpc_translated_chunks = [] for chunk in chunks: if "Original" in translation_methods: translation, _ = original_translation(chunk, src_language, target_language, session_id) org_translated_chunks.append(translation) if "Plan2Align" in translation_methods: translation, _ = plan2align_translate_text( chunk, session_id, model, tokenizer, device, src_language, target_language, max_iterations_value, threshold_value, good_ref_contexts_num_value, "metricx" ) p2a_translated_chunks.append(translation) if "Best-of-N" in translation_methods: translation, _ = best_of_n_translation(chunk, src_language, target_language, max_iterations_value, session_id) bfn_translated_chunks.append(translation) if "MPC" in translation_methods: translation, _ = mpc_translation(chunk, src_language, target_language, max_iterations_value, session_id) mpc_translated_chunks.append(translation) org_combined_translation = ' '.join(org_translated_chunks) p2a_combined_translation = ' '.join(p2a_translated_chunks) bfn_combined_translation = ' '.join(bfn_translated_chunks) mpc_combined_translation = ' '.join(mpc_translated_chunks) orig, best_score = summary_translate(text, org_combined_translation, target_language, session_id) orig_output = f"{orig}\n\nScore: {best_score:.2f}" plan2align_trans, best_score = summary_translate(text, p2a_combined_translation, target_language, session_id) plan2align_output = f"{plan2align_trans}\n\nScore: {best_score:.2f}" best_candidate, best_score = summary_translate(text, bfn_combined_translation, target_language, session_id) best_of_n_output = f"{best_candidate}\n\nScore: {best_score:.2f}" mpc_candidate, best_score = summary_translate(text, mpc_combined_translation, target_language, session_id) mpc_output = f"{mpc_candidate}\n\nScore: {best_score:.2f}" return orig_output, plan2align_output, best_of_n_output, mpc_output # ---------- Gradio ---------- target_languages = ["Traditional Chinese", "Simplified Chinese", "English", "Russian", "German", "Japanese", "Korean"] src_languages = ["Traditional Chinese", "Simplified Chinese", "English", "Russian", "German", "Japanese", "Korean"] with gr.Blocks(title="Test-Time Machine Translation with Plan2Align") as demo: state = gr.State({}) gr.Markdown("# Translation Demo: Multiple Translation Methods") gr.Markdown("請選擇要執行的翻譯方法（可多選或全選）：") with gr.Row(): with gr.Column(scale=1): source_text = gr.Textbox( label="Source Text", placeholder="請輸入文本...", lines=5 ) src_language_input = gr.Dropdown( choices=src_languages, value="Traditional Chinese", label="Source Language" ) task_language_input = gr.Dropdown( choices=target_languages, value="English", label="Target Language" ) max_iterations_input = gr.Number(label="Max Iterations", value=6) threshold_input = gr.Number(label="Threshold", value=0.7) good_ref_contexts_num_input = gr.Number(label="Good Ref Contexts Num", value=5) translation_methods_input = gr.CheckboxGroup( choices=["Original", "Plan2Align", "Best-of-N", "MPC"], value=["Original", "Plan2Align"], label="Translation Methods" ) chunk_size_input = gr.Number( # ✅ add chunk function label="Chunk Size (-1 for all)", value=-1 ) translate_button = gr.Button("Translate") with gr.Column(scale=2): original_output = gr.Textbox( label="Original Translation", lines=5, interactive=False ) plan2align_output = gr.Textbox( label="Plan2Align Translation", lines=5, interactive=False ) best_of_n_output = gr.Textbox( label="Best-of-N Translation", lines=5, interactive=False ) mpc_output = gr.Textbox( label="MPC Translation", lines=5, interactive=False ) translate_button.click( fn=process_text, inputs=[ source_text, src_language_input, task_language_input, max_iterations_input, threshold_input, good_ref_contexts_num_input, translation_methods_input, chunk_size_input, # ✅ add chunk function state ], outputs=[original_output, plan2align_output, best_of_n_output, mpc_output] ) gr.Examples( examples=[ ["台灣夜市文化豐富多彩，從士林夜市到饒河街夜市，提供各種美食、遊戲和購物體驗，吸引了無數遊客。", "Traditional Chinese", "English", 2, 0.7, 1, ["Original", "Plan2Align"], -1], ["台北101曾經是世界最高的建築物，它不僅是台灣的地標，也象徵著經濟成就和創新精神。", "Traditional Chinese", "Japanese", 2, 0.7, 1, ["Original", "Plan2Align"], -1], ["阿里山日出和森林鐵路是台灣最著名的自然景觀之一，每年吸引數十萬遊客前來欣賞雲海和壯麗的日出。", "Traditional Chinese", "Korean", 2, 0.7, 1, ["Original", "Plan2Align"], -1], # ["珍珠奶茶，這款源自台灣的獨特飲品，不僅在台灣本地深受喜愛，更以其獨特的風味和口感，在全球掀起了一股熱潮，成為了一種跨越文化、風靡全球的時尚飲品。", "Traditional Chinese", "Japanese", 3, 0.7, 3, ["Original", "Plan2Align"], -1], # ["原住民文化如同一片深邃的星空，閃爍著無數璀璨的傳統與藝術光芒。他們的歌舞，是與祖靈對話的旋律，是與自然共鳴的節奏，每一個舞步、每一聲吟唱，都承載著古老的傳說與智慧。編織，是他們巧手下的詩篇，一絲一線，交織出生命的紋理，也編織出對土地的熱愛與敬畏。木雕，則是他們與自然對話的雕塑，每一刀、每一鑿，都刻畫著對萬物的觀察與敬意，也雕琢出對祖先的追憶與傳承。", "Traditional Chinese", "Korean", 5, 0.7, 5, ["Original", "Plan2Align"], -1] ], inputs=[ source_text, src_language_input, task_language_input, max_iterations_input, threshold_input, good_ref_contexts_num_input, translation_methods_input, chunk_size_input # ✅ add chunk function ], outputs=[original_output, plan2align_output, best_of_n_output, mpc_output], fn=process_text ) gr.Markdown("## How It Works") gr.Markdown(""" 1. **Original Translation:** 利用固定提示生成候選，直接取首個候選作為原始翻譯。 2. **Plan2Align Translation:** 採用 context alignment 和 self-rewriting 策略進行翻譯，適合長文翻譯。 3. **Best-of-N Translation:** 重複生成多次候選，評分選出最佳翻譯，適合短文翻譯。 4. **MPC Translation:** 以迭代改善策略，每輪生成候選後評分，並將最佳翻譯作為下一輪輸入，適合短文翻譯。 若輸入文本超過 1024 tokens，Best-of-N 與 MPC 方法會回傳警告訊息。 """) if __name__ == "__main__": demo.launch(share=True)