[test] demo

app.py

@@ -1,154 +1,214 @@
 import gradio as gr
-import numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
 import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from trl import AutoModelForCausalLMWithValueHead
+
+# Set device and dtype
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.bfloat16
+
+# Load models only once at startup
+print("Loading models...")
+model_id = "meta-llama/Meta-Llama-3.1-8B"  # Replace with your actual model ID
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+lm_model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch_dtype,
+    device_map="auto"
+)
+
+# Load the reward model
+RM = AutoModelForCausalLMWithValueHead.from_pretrained(
+    'ray24724919/plan2align_rm',
+    torch_dtype=torch_dtype,
+    device_map="auto"
+)
+RM.eval()
+print("Models loaded successfully!")
+
+# Self-contained translation and evaluation functions
+def translate(source_text, target_language="English"):
+    """
+    Translate text from Chinese to the specified target language.
+    
+    Args:
+        source_text (str): The Chinese text to translate
+        target_language (str): The target language for translation
+        
+    Returns:
+        str: The translated text
+    """
+    # Format the input as per the system prompt
+    messages = [
+        {"role": "system", "content": "You are a helpful translator and only output the result."},
+        {"role": "user", "content": f"### Translate this from Chinese to {target_language}, Chinese:\n{source_text}\n### {target_language}:"}
+    ]
+    
+    # Format messages for the model
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Generate translation
+    with torch.no_grad():
+        outputs = lm_model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode the generated text
+    translation = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
+    return translation
+
+def evaluate_translation(source_text, translation, target_language="English"):
+    """
+    Evaluate the quality of a translation using the reward model.
+    
+    Args:
+        source_text (str): The original Chinese text
+        translation (str): The translated text
+        target_language (str): The target language of the translation
+        
+    Returns:
+        float: The reward score
+    """
+    messages = [
+        {"role": "system", "content": "You are a helpful translator and only output the result."},
+        {"role": "user", "content": f"### Translate this from Chinese to {target_language}, Chinese:\n{source_text}\n### {target_language}:"},
+        {"role": "assistant", "content": translation}
+    ]
+    
+    # Format messages for the reward model
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+    
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Get reward score
+    with torch.no_grad():
+        outputs = RM(input_ids=inputs.input_ids)
+        reward_score = outputs.value.item()
+    
+    return reward_score
+
+# Combined function for the Gradio interface
+def translate_text(source_text, target_language):
+    """
+    Translate text and get reward score
+    
+    Args:
+        source_text (str): The Chinese text to translate
+        target_language (str): The target language for translation
+        
+    Returns:
+        tuple: (translation, reward_score)
+    """
+    if not source_text.strip():
+        return "Please enter some text to translate.", 0.0
+    
+    try:
+        translation = translate(source_text, target_language)
+        reward_score = evaluate_translation(source_text, translation, target_language)
+        return translation, float(reward_score)
+    except Exception as e:
+        return f"Error: {str(e)}", 0.0
+
+# Define available target languages
+target_languages = [
+    "English", "French", "Spanish", "German", "Italian", 
+    "Portuguese", "Russian", "Japanese", "Korean", "Arabic"
 ]
 
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
+# Create the Gradio interface
+with gr.Blocks(title="Chinese Translation with Reward Scoring") as demo:
+    gr.Markdown("# Chinese to Any Language Translation")
+    gr.Markdown("This demo translates Chinese text to your chosen language and provides a quality score from our reward model.")
+    
+    with gr.Row():
+        with gr.Column():
+            source_text = gr.Textbox(
+                label="Chinese Text",
+                placeholder="Enter Chinese text here...",
+                lines=5
             )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+            target_language = gr.Dropdown(
+                choices=target_languages,
+                value="English",
+                label="Target Language"
             )
-
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+            translate_button = gr.Button("Translate")
+        
+        with gr.Column():
+            translation_output = gr.Textbox(
+                label="Translation",
+                lines=5,
+                interactive=False
             )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
+            reward_score = gr.Number(
+                label="Translation Quality Score (higher is better)",
+                precision=4,
+                interactive=False
+            )
+            
             with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
+                score_indicator = gr.Label(label="Quality Rating")
+    
+    # Function to update the quality rating based on score
+    def update_quality_rating(score):
+        if score >= 0.8:
+            return "Excellent"
+        elif score >= 0.6:
+            return "Good"
+        elif score >= 0.4:
+            return "Average"
+        elif score >= 0.2:
+            return "Poor"
+        else:
+            return "Very Poor"
+    
+    # Set up the translation flow
+    translate_outputs = translate_button.click(
+        fn=translate_text,
+        inputs=[source_text, target_language],
+        outputs=[translation_output, reward_score]
+    )
+    
+    # Update the quality rating whenever the reward score changes
+    reward_score.change(
+        fn=update_quality_rating,
+        inputs=[reward_score],
+        outputs=[score_indicator]
+    )
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            ["你好，世界！", "English"],
+            ["我喜欢学习新的语言。", "Spanish"],
+            ["北京烤鴨很好吃。", "French"],
+            ["人工智能正在改变世界。", "German"],
+            ["今天天气真好。", "Japanese"]
         ],
-        outputs=[result, seed],
+        inputs=[source_text, target_language],
+        outputs=[translation_output, reward_score],
+        fn=translate_text
     )
-
+    
+    gr.Markdown("## How It Works")
+    gr.Markdown("""
+    1. Enter Chinese text in the input box
+    2. Select your desired target language
+    3. Click 'Translate' to get the translation
+    4. The system will display the translation and a quality score
+    
+    The quality score is generated by a reward model trained to evaluate translation quality.
+    Higher scores indicate better translations.
+    """)
+
+# Launch the app
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1,6 +1,7 @@
 accelerate
-diffusers
-invisible_watermark
-torch
-transformers
+gradio
+safetensors
+torch>=2.0.0
+transformers>=4.30.0
+trl>=0.7.1
 xformers

translation_model.py

@@ -0,0 +1,158 @@
+import torch
+import safetensors.torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from trl import AutoModelForCausalLMWithValueHead
+
+# Set device and dtype
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.bfloat16
+
+# Load the base LLaMa 3.1 8B model for translation
+model_id = "meta-llama/Meta-Llama-3.1-8B"  # Replace with your actual model ID
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+lm_model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch_dtype,
+    device_map="auto"
+)
+
+# Load the reward model
+RM = AutoModelForCausalLMWithValueHead.from_pretrained(
+    'ray24724919/plan2align_rm',
+    torch_dtype=torch_dtype,
+    device_map="auto"
+)
+RM.eval()
+RM.gradient_checkpointing_enable()  # if needed for memory efficiency
+
+# Define the load_file function
+def load_file(file_path):
+    return safetensors.torch.load_file(file_path)
+
+# Load value head weights if you have the file
+# If you don't have the specific file, you might need to download it or use the model as is
+try:
+    value_head_weights = load_file("value_head.safetensors")  # Replace with actual path
+    new_state_dict = {key.replace("v_head.", "") if key.startswith("v_head.") else key: value for key, value in value_head_weights.items()}
+    RM.v_head.load_state_dict(new_state_dict)
+except FileNotFoundError:
+    print("Value head weights file not found. Using default weights.")
+
+# Define translation function with more flexibility
+def translate(source_text, target_language="English", model=lm_model):
+    """
+    Translate text from Chinese to the specified target language.
+    
+    Args:
+        source_text (str): The Chinese text to translate
+        target_language (str): The target language for translation
+        model: The model to use for translation
+        
+    Returns:
+        str: The translated text
+    """
+    # Format the input as per the system prompt
+    messages = [
+        {"role": "system", "content": "You are a helpful translator and only output the result."},
+        {"role": "user", "content": f"### Translate this from Chinese to {target_language}, Chinese:\n{source_text}\n### {target_language}:"}
+    ]
+    
+    # Format messages for the model
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Generate translation
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode the generated text
+    translation = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
+    return translation
+
+# Evaluate the translation using the reward model
+def evaluate_translation(source_text, translation, target_language="English"):
+    """
+    Evaluate the quality of a translation using the reward model.
+    
+    Args:
+        source_text (str): The original Chinese text
+        translation (str): The translated text
+        target_language (str): The target language of the translation
+        
+    Returns:
+        float: The reward score
+    """
+    messages = [
+        {"role": "system", "content": "You are a helpful translator and only output the result."},
+        {"role": "user", "content": f"### Translate this from Chinese to {target_language}, Chinese:\n{source_text}\n### {target_language}:"},
+        {"role": "assistant", "content": translation}
+    ]
+    
+    # Format messages for the reward model
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+    
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Get reward score
+    with torch.no_grad():
+        outputs = RM(input_ids=inputs.input_ids)
+        reward_score = outputs.value.item()
+    
+    return reward_score
+
+# Function to translate and evaluate in one step
+def translate_and_evaluate(source_text, target_language="English"):
+    """
+    Translate text and evaluate the translation quality in one step.
+    
+    Args:
+        source_text (str): The Chinese text to translate
+        target_language (str): The target language for translation
+        
+    Returns:
+        tuple: (translation, reward_score)
+    """
+    translation = translate(source_text, target_language)
+    reward_score = evaluate_translation(source_text, translation, target_language)
+    return translation, reward_score
+
+# Example usage
+if __name__ == "__main__":
+    # Example with default target language (English)
+    source = "你好世界"
+    translation, reward_score = translate_and_evaluate(source)
+    print(f"Source: {source}")
+    print(f"Translation to English: {translation}")
+    print(f"Reward Score: {reward_score}")
+    
+    # Example with custom target language
+    target_language = "French"
+    translation, reward_score = translate_and_evaluate(source, target_language)
+    print(f"\nSource: {source}")
+    print(f"Translation to {target_language}: {translation}")
+    print(f"Reward Score: {reward_score}")
+    
+    # Interactive mode
+    print("\n=== Interactive Translation Mode ===")
+    print("Enter 'quit' to exit")
+    while True:
+        user_input = input("\nEnter Chinese text to translate: ")
+        if user_input.lower() == 'quit':
+            break
+            
+        target = input("Enter target language (default: English): ").strip()
+        if not target:
+            target = "English"
+            
+        translation, reward_score = translate_and_evaluate(user_input, target)
+        print(f"Translation to {target}: {translation}")
+        print(f"Reward Score: {reward_score}")