polish the demo

app.py

@@ -1,4 +1,5 @@
 import os
+import gc
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -124,6 +125,14 @@ class RewardModel:
 reward_model = RewardModel(device, tokenizer, torch_dtype=torch.float16)
 print("Models loaded successfully!")
 
+# Memory management function
+def clear_cache():
+    """Clear CUDA cache and run garbage collection to free memory"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+    return "Cache cleared"
+
 # Helper functions from Plan2Align
 def rm_predict_preference(source, translation0, translation1, language="English"):
     translations = [translation0, translation1]
@@ -202,7 +211,7 @@ def translate_chinese_to_english(chinese_text, target_language="English"):
         "best_index": best_index
     }
 
-# Updated Gradio interface
+# Updated process_text function with cache clearing
 def process_text(text, target_language="English"):
     if not text.strip():
         return "Please enter some text to translate.", "", "", "", ""
@@ -218,6 +227,9 @@ def process_text(text, target_language="English"):
         
         candidates_text = "\n".join(candidates)
         
+        # Clear cache after processing
+        clear_cache()
+        
         return (
             result["best_translation"],
             f"{result['best_reward']:.4f}",
@@ -226,6 +238,8 @@ def process_text(text, target_language="English"):
             "Yes" if result["best_reward"] >= THRESHOLD else "No"
         )
     except Exception as e:
+        # Clear cache even if there's an error
+        clear_cache()
         return f"Error: {str(e)}", "", "", "", ""
 
 # Define available target languages - only the supported ones
@@ -234,9 +248,8 @@ target_languages = [
 ]
 
 # 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.Blocks(title="Test-Time Machine 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. Paper: https://arxiv.org/pdf/2502.20795.")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -251,6 +264,7 @@ with gr.Blocks(title="Chinese Translation with Plan2Align") as demo:
                 label="Target Language"
             )
             translate_button = gr.Button("Translate")
+            clear_button = gr.Button("Clear Memory Cache")
         
         with gr.Column(scale=2):
             with gr.Tab("Best Translation"):
@@ -278,6 +292,12 @@ with gr.Blocks(title="Chinese Translation with Plan2Align") as demo:
                     lines=15,
                     interactive=False
                 )
+            
+            cache_status = gr.Textbox(
+                label="Cache Status",
+                value="Ready",
+                interactive=False
+            )
     
     # Set up the translation flow
     translate_button.click(
@@ -286,7 +306,14 @@ with gr.Blocks(title="Chinese Translation with Plan2Align") as demo:
         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
+    # Add manual cache clearing button
+    clear_button.click(
+        fn=clear_cache,
+        inputs=[],
+        outputs=[cache_status]
+    )
+    
+    # Examples with more complex sentences in Traditional Chinese about Taiwan
     gr.Examples(
         examples=[
             ["夜市文化豐富多彩，從士林夜市到饒河街夜市，提供各種美食、遊戲和購物體驗，吸引了無數遊客。", "English"],