polish the demo

app.py

@@ -155,7 +155,7 @@ def rm_find_best_translation(source, translations, language="English"):
     else:
         return None
 
-def translate_chinese_to_english(chinese_text):
+def translate_chinese_to_english(chinese_text, target_language="English"):
     # Generate multiple translations
     translations = []
     
@@ -169,7 +169,7 @@ def translate_chinese_to_english(chinese_text):
     for prompt in system_prompts:
         messages = [
             {"role": "system", "content": prompt},
-            {"role": "user", "content": f"Translate the following Chinese text to English:\n\n{chinese_text}"}
+            {"role": "user", "content": f"Translate the following Chinese text to {target_language}:\n\n{chinese_text}"}
         ]
         
         inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(device)
@@ -185,26 +185,134 @@ def translate_chinese_to_english(chinese_text):
         translation = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
         translations.append(translation)
     
-    # Use reward model to find the best translation
-    best_translation = rm_find_best_translation(chinese_text, translations)
+    # Get rewards for all translations
+    rewards = reward_model.reward_fn(target_language, chinese_text.replace('</s>',' '), 
+                                    [t.replace('</s>',' ') for t in translations])
     
-    if best_translation is None:
-        # If no translation meets the threshold, return the first one
-        return translations[0]
+    # Find the best translation
+    best_index = rewards.index(max(rewards))
+    best_translation = translations[best_index]
+    
+    # Return all information
+    return {
+        "best_translation": best_translation,
+        "best_reward": rewards[best_index],
+        "all_translations": translations,
+        "all_rewards": rewards,
+        "best_index": best_index
+    }
+
+# Updated Gradio interface
+def process_text(text, target_language="English"):
+    if not text.strip():
+        return "Please enter some text to translate.", "", "", "", ""
     
-    return best_translation
+    try:
+        result = translate_chinese_to_english(text, target_language)
+        
+        # Format the candidate translations with their rewards
+        candidates = []
+        for i, (trans, reward) in enumerate(zip(result["all_translations"], result["all_rewards"])):
+            marker = "★ " if i == result["best_index"] else ""
+            candidates.append(f"{marker}Candidate {i+1} (Reward: {reward:.4f}):\n{trans}\n")
+        
+        candidates_text = "\n".join(candidates)
+        
+        return (
+            result["best_translation"],
+            f"{result['best_reward']:.4f}",
+            candidates_text,
+            f"Candidate {result['best_index']+1}",
+            "Yes" if result["best_reward"] >= THRESHOLD else "No"
+        )
+    except Exception as e:
+        return f"Error: {str(e)}", "", "", "", ""
 
-# Gradio interface
-def process_text(text):
-    return translate_chinese_to_english(text)
+# Define available target languages - only the supported ones
+target_languages = [
+    "English", "Russian", "German", "Japanese", "Korean"
+]
 
-demo = gr.Interface(
-    fn=process_text,
-    inputs=gr.Textbox(lines=5, placeholder="Enter Chinese text here..."),
-    outputs=gr.Textbox(lines=5),
-    title="Chinese to English Translation with Plan2Align",
-    description="This app uses the Plan2Align approach to translate Chinese text to English."
-)
+# Create an enhanced Gradio interface
+with gr.Blocks(title="Chinese Translation with Plan2Align") as demo:
+    gr.Markdown("# Chinese Translation with Plan2Align")
+    gr.Markdown("This demo uses the Plan2Align approach to translate Chinese text to your chosen language, showing how the reward model evaluates different translation candidates.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            source_text = gr.Textbox(
+                label="Chinese Text",
+                placeholder="Enter Chinese text here...",
+                lines=5
+            )
+            target_language = gr.Dropdown(
+                choices=target_languages,
+                value="English",
+                label="Target Language"
+            )
+            translate_button = gr.Button("Translate")
+        
+        with gr.Column(scale=2):
+            with gr.Tab("Best Translation"):
+                best_translation = gr.Textbox(
+                    label="Best Translation",
+                    lines=5,
+                    interactive=False
+                )
+                best_reward = gr.Textbox(
+                    label="Reward Score",
+                    interactive=False
+                )
+                best_candidate = gr.Textbox(
+                    label="Best Candidate",
+                    interactive=False
+                )
+                meets_threshold = gr.Textbox(
+                    label="Meets Quality Threshold",
+                    interactive=False
+                )
+            
+            with gr.Tab("All Candidates"):
+                all_candidates = gr.Textbox(
+                    label="All Translation Candidates with Rewards",
+                    lines=15,
+                    interactive=False
+                )
+    
+    # Set up the translation flow
+    translate_button.click(
+        fn=process_text,
+        inputs=[source_text, target_language],
+        outputs=[best_translation, best_reward, all_candidates, best_candidate, meets_threshold]
+    )
+    
+    # Examples with more complex sentences in Traditional Chinese about Taiwan for the supported languages
+    gr.Examples(
+        examples=[
+            ["夜市文化豐富多彩，從士林夜市到饒河街夜市，提供各種美食、遊戲和購物體驗，吸引了無數遊客。", "English"],
+            ["台北101曾經是世界最高的建築物，它不僅是台灣的地標，也象徵著經濟成就和創新精神。", "Russian"],
+            ["阿里山日出和森林鐵路是台灣最著名的自然景觀之一，每年吸引數十萬遊客前來欣賞雲海和壯麗的日出。", "German"],
+            ["珍珠奶茶起源於台灣，現已成為全球流行的飲品，展現了飲食文化對世界的影響力。", "Japanese"],
+            ["原住民文化擁有豐富的傳統和藝術表現形式，包括歌舞、編織和木雕，反映了與自然和諧共處的生活智慧。", "Korean"]
+        ],
+        inputs=[source_text, target_language],
+        outputs=[best_translation, best_reward, all_candidates, best_candidate, meets_threshold],
+        fn=process_text
+    )
+    
+    gr.Markdown("## How It Works")
+    gr.Markdown("""
+    1. The system generates three different translations using different translation styles:
+       - Literal: A word-for-word translation preserving structure
+       - Professional: A clear, formal translation
+       - Creative: A vivid, expressive translation
+    
+    2. The reward model evaluates each translation and assigns a score
+    
+    3. The translation with the highest reward score is selected as the best
+    
+    4. A translation meets the quality threshold if its reward score is ≥ 2.0
+    """)
 
 if __name__ == "__main__":
     demo.launch()