add chunk function

app.py

@@ -69,21 +69,33 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16
 )
 
-# def generate_translation(system_prompt, prompt):
-#     messages=[
-#         {"role": "system", "content": system_prompt},
-#         {"role": "user", "content": prompt}
-#     ]
-#     inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
-#     outputs = model.generate(
-#             inputs,
-#             max_new_tokens=512,
-#             temperature=0.7,
-#             top_p=0.9,
-#             do_sample=True
-#         )
-#     translation = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
-#     return translation
+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:"
@@ -160,6 +172,21 @@ def basic_translate(source_sentence, src_language, tgt_language):
         translations.append(translation)
     return translations 
 
+def summary_translate(src_text, temp_tgt_text, tgt_language):
+    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 {textemp_tgt_textt}"
+        prompt += f"\n### Rephrased:\n"
+        translation = generate_translation(prompt_style, prompt)
+        translations.append(translation)
+
+    best, score = evaluate_candidates(src_text, translations, target_language, session_id)
+    if cand_list:
+        return best, score
+    return "", 0
+
 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,
@@ -255,23 +282,67 @@ def process_text(text, src_language, target_language, max_iterations_value, thre
     best_of_n_output = ""
     mpc_output = ""
 
-
-    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}"
+    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:
+        chunks = [' '.join(source_segments)]
+    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(org_combined_translation, target_language)
+    orig_output = f"{orig}\n\nScore: {best_score:.2f}"
+    plan2align_trans, best_score = summary_translate(p2a_combined_translation, target_language)
+    plan2align_output = f"{plan2align_trans}\n\nScore: {best_score:.2f}"
+    best_candidate, best_score = summary_translate(bfn_combined_translation, target_language)
+    best_of_n_output = f"{best_candidate}\n\nScore: {best_score:.2f}"
+    mpc_candidate, best_score = summary_translate(mpc_combined_translation, target_language)
+    mpc_output = f"{mpc_candidate}\n\nScore: {mpc_score:.2f}"
+
+    # 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}"
     
     return orig_output, plan2align_output, best_of_n_output, mpc_output
 
@@ -310,6 +381,10 @@ with gr.Blocks(title="Test-Time Machine Translation with Plan2Align") as demo:
                 value=["Original", "Plan2Align"],
                 label="Translation Methods"
             )
+            chunk_size_input = gr.Number(  # ✅ add chunk function
+                label="Chunk Size (Number of sentences per translation, -1 for all)",
+                value=-1
+            )
             translate_button = gr.Button("Translate")
         with gr.Column(scale=2):
             original_output = gr.Textbox(
@@ -343,6 +418,7 @@ with gr.Blocks(title="Test-Time Machine Translation with Plan2Align") as demo:
             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]
@@ -350,11 +426,11 @@ with gr.Blocks(title="Test-Time Machine Translation with Plan2Align") as demo:
     
     gr.Examples(
         examples=[
-            ["台灣夜市文化豐富多彩，從士林夜市到饒河街夜市，提供各種美食、遊戲和購物體驗，吸引了無數遊客。", "Traditional Chinese", "English", 2, 0.7, 1, ["Original", "Plan2Align"]],
-            ["台北101曾經是世界最高的建築物，它不僅是台灣的地標，也象徵著經濟成就和創新精神。", "Traditional Chinese", "Russian", 2, 0.7, 1, ["Original", "Plan2Align"]],
-            ["阿里山日出和森林鐵路是台灣最著名的自然景觀之一，每年吸引數十萬遊客前來欣賞雲海和壯麗的日出。", "Traditional Chinese", "German", 2, 0.7, 1, ["Original", "Plan2Align"]],
-            ["珍珠奶茶，這款源自台灣的獨特飲品，不僅在台灣本地深受喜愛，更以其獨特的風味和口感，在全球掀起了一股熱潮，成為了一種跨越文化、風靡全球的時尚飲品。", "Traditional Chinese", "Japanese", 3, 0.7, 3, ["Original", "Plan2Align"]],
-            ["原住民文化如同一片深邃的星空，閃爍著無數璀璨的傳統與藝術光芒。他們的歌舞，是與祖靈對話的旋律，是與自然共鳴的節奏，每一個舞步、每一聲吟唱，都承載著古老的傳說與智慧。編織，是他們巧手下的詩篇，一絲一線，交織出生命的紋理，也編織出對土地的熱愛與敬畏。木雕，則是他們與自然對話的雕塑，每一刀、每一鑿，都刻畫著對萬物的觀察與敬意，也雕琢出對祖先的追憶與傳承。", "Traditional Chinese", "Korean", 5, 0.7, 5, ["Original", "Plan2Align"]]
+            ["台灣夜市文化豐富多彩，從士林夜市到饒河街夜市，提供各種美食、遊戲和購物體驗，吸引了無數遊客。", "Traditional Chinese", "English", 2, 0.7, 1, ["Original", "Plan2Align"], -1],
+            ["台北101曾經是世界最高的建築物，它不僅是台灣的地標，也象徵著經濟成就和創新精神。", "Traditional Chinese", "Russian", 2, 0.7, 1, ["Original", "Plan2Align"], -1],
+            ["阿里山日出和森林鐵路是台灣最著名的自然景觀之一，每年吸引數十萬遊客前來欣賞雲海和壯麗的日出。", "Traditional Chinese", "German", 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,
@@ -363,7 +439,8 @@ with gr.Blocks(title="Test-Time Machine Translation with Plan2Align") as demo:
             max_iterations_input,
             threshold_input,
             good_ref_contexts_num_input,
-            translation_methods_input
+            translation_methods_input,
+            chunk_size_input  # ✅ add chunk function
         ],
         outputs=[original_output, plan2align_output, best_of_n_output, mpc_output],
         fn=process_text