Spaces:
Sleeping
Sleeping
| 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) | |