Update app.py

app.py

@@ -1,64 +1,377 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import os
+import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+import re
+import torch
 
+# Set environment variables
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+DESCRIPTION = '''
+<div>
+<h1 style="text-align: center;">JudgeLRM</h1>
+<p>This Space demonstrates the <a href="https://huggingface.co/nuojohnchen/JudgeLRM-7B"><b>JudgeLRM</b></a> model, designed to evaluate the quality of two AI assistant responses. JudgeLRM is a family of judgment-oriented LLMs trained using reinforcement learning (RL) with judge-wise, outcome-driven rewards. JudgeLRM models consistently outperform both SFT-tuned and state-of-the-art reasoning models. Notably, JudgeLRM-3B surpasses GPT-4, and JudgeLRM-7B outperforms DeepSeek-R1 by 2.79\% in F1 score, particularly excelling in judge tasks requiring deep reasoning.</p>
+<p>Enter an instruction and two responses, and the model will score them on a scale of 1-10 (higher is better).</p>
+<p>You can also select Hugging Face models to automatically generate responses for evaluation.</p>
+</div>
+'''
+
+LICENSE = """
+<div style="font-family: monospace; white-space: pre; margin-top: 20px; line-height: 1.2;">
+@misc{XtraGPT,
+    title = {JudgeLRM},
+    url = {https://huggingface.co/nuojohnchen/JudgeLRM-7B},
+    author = {Nuo Chen, Zhiyuan Hu, Qingyun Zou, Jiaying Wu, Qian Wang, Bryan Hooi, Bingsheng He},
+    month = {March},
+    year = {2025}
+}
+</div>
 """
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
+
+PLACEHOLDER = """
+<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
+   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">JudgeLRM</h1>
+   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Enter an instruction and two responses, I'll evaluate their quality...</p>
+</div>
 """
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+css = """
+h1 {
+  text-align: center;
+  display: block;
+}
+#duplicate-button {
+  margin: auto;
+  color: white;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+"""
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+# Model paths
+MODEL_PATHS = {
+    "JudgeLRM-3B": "nuojohnchen/JudgeLRM-3B",
+    "JudgeLRM-7B": "nuojohnchen/JudgeLRM-7B"
+}
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+# Popular models for dropdown selection
+POPULAR_MODELS = [
+    "Qwen/Qwen2.5-7B-Instruct",
+    "01-ai/Yi-6B-Chat",
+    "microsoft/phi-2",
+    "FreedomIntelligence/Apollo-7B",
+    "tiiuae/falcon-7b-instruct",
+    "HuggingFaceH4/zephyr-7b-beta",
+    "stabilityai/stablelm-3b-4e1t",
+    "openchat/openchat-3.5-0106"
+]
 
-    messages.append({"role": "user", "content": message})
+# Global variables for model and tokenizer
+tokenizer = None
+model = None
+current_model_name = None
+response_model_1 = None
+response_model_2 = None
+response_tokenizer_1 = None
+response_tokenizer_2 = None
 
-    response = ""
+def extract_scores(text):
+    """Extract scores from generated text"""
+    pattern = r'<answer>(\d+)</answer><answer>(\d+)</answer>'
+    match = re.search(pattern, text)
+    if match:
+        return int(match.group(1)), int(match.group(2))
+    return None, None
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+# Function to determine which model path to use
+def get_model_path(dropdown_value, custom_value):
+    """Return custom value if provided, otherwise return dropdown value"""
+    if custom_value and custom_value.strip():
+        return custom_value.strip()
+    return dropdown_value
 
-        response += token
-        yield response
+# Function to generate response from a model
+def generate_response(instruction, model_path, progress=gr.Progress()):
+    """Generate a response from a specified model"""
+    progress(0, desc=f"Loading model {model_path}...")
+    try:
+        # Load model and tokenizer
+        response_tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        response_model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.float16)
+        
+        progress(0.5, desc=f"Generating response from {model_path}...")
+        
+        # Create prompt (simple format, adjust as needed for specific models)
+        prompt = f"<|user|>\n{instruction}\n<|assistant|>"
+        
+        # Generate response
+        input_ids = response_tokenizer.encode(prompt, return_tensors="pt").to(response_model.device)
+        output = response_model.generate(
+            input_ids=input_ids,
+            max_new_tokens=1024,
+            temperature=0.7,
+            do_sample=True
+        )
+        
+        # Decode response and clean it
+        full_response = response_tokenizer.decode(output[0], skip_special_tokens=True)
+        
+        # Remove the prompt part from the response
+        clean_response = full_response.replace(f"<|user|>\n{instruction}\n<|assistant|>", "").strip()
+        
+        # If the model doesn't use these exact tokens, try to extract just the assistant's response
+        if clean_response == full_response:
+            # Try to find where the assistant's response starts
+            parts = full_response.split(instruction)
+            if len(parts) > 1:
+                clean_response = parts[1].strip()
+                # Further clean any remaining tokens
+                for token in ["<|assistant|>", "<assistant>", "Assistant:", "A:"]:
+                    clean_response = clean_response.replace(token, "").strip()
+        
+        # Clean up resources
+        del response_model
+        del response_tokenizer
+        torch.cuda.empty_cache()
+        
+        progress(1.0, desc=f"Response from {model_path} generated")
+        return clean_response
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
+@spaces.GPU(duration=120)
+def judge_responses(instruction, response1, response2, model_name, temperature=0.1, max_new_tokens=2048):
+    """
+    Evaluate the quality of two responses
+    Args:
+        instruction (str): Instruction/question
+        response1 (str): First response
+        response2 (str): Second response
+        model_name (str): Model to use for evaluation
+        temperature (float): Generation temperature
+        max_new_tokens (int): Maximum number of tokens to generate
+    Returns:
+        str: Generated evaluation result
+    """
+    global tokenizer, model, current_model_name
+    
+    # Load model on demand if it's not already loaded or if a different model is requested
+    if model is None or model_name != current_model_name:
+        # Clear GPU memory if a model is already loaded
+        if model is not None:
+            del model
+            del tokenizer
+            torch.cuda.empty_cache()
+        
+        # Update status
+        yield f"Loading {model_name}... Please wait."
+        
+        # Load the requested model
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATHS[model_name], use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(MODEL_PATHS[model_name], device_map="auto")
+            current_model_name = model_name
+        except Exception as e:
+            yield f"Error loading model: {str(e)}"
+            return
+    
+    # Build prompt
+    prompt = """<|im_start|>system\nYou are a helpful assistant. The assistant first performs a detailed, step-by-step reasoning process in its mind and then provides the user with the answer. The reasoning process and answer are enclosed within <think> </think> and<answer> </answer> tags, respectively, i.e., <think> detailed reasoning process here, explaining each step of your evaluation for both assistants </think><answer> answer here </answer>. Now the user asks you to judge the performance of two AI assistants in response to the question. Score assistants 1-10 (higher=better). Criteria includes helpfulness, relevance, accuracy, and level of detail. Avoid order, length, style or other bias. After thinking, when you finally reach a conclusion, clearly  provide your evaluation scores within <answer> </answer> tags, i.e. for example,<answer>3</answer><answer>5</answer>\n<|im_end|>\n<|im_start|>user\n[Question]\n{question}\n\n[Assistant 1's Answer]\n{answer_1}\n\n[Assistant 2's Answer]\n{answer_2}\n<|im_end|>\n<|im_start|>assistant\n"""
+    
+    formatted_prompt = prompt.format(question=instruction, answer_1=response1, answer_2=response2)
+    
+    # Set up streaming output
+    input_ids = tokenizer.encode(formatted_prompt, return_tensors="pt").to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=False)
+    
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+    )
+    
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+    
+    # Run generation in a separate thread
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    
+    # Collect generated text
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        full_text = "".join(outputs)
+        
+        # Try to extract scores
+        score1, score2 = extract_scores(full_text)
+        if score1 and score2:
+            result = f"{full_text}\n\n**Evaluation Results:** Response 1 Score: {score1}/10, Response 2 Score: {score2}/10"
+        else:
+            result = full_text
+            
+        yield result
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+@spaces.GPU(duration=120)
+def generate_and_judge(instruction, model_dropdown_1, custom_model_1, model_dropdown_2, custom_model_2, judge_model_name, temperature=0.1, max_new_tokens=2048, progress=gr.Progress()):
+    """Generate responses from two models and judge them"""
+    progress(0, desc="Starting generation process")
+    
+    # Determine which model paths to use
+    model_path_1 = get_model_path(model_dropdown_1, custom_model_1)
+    model_path_2 = get_model_path(model_dropdown_2, custom_model_2)
+    
+    # Generate responses from both models
+    progress(0.1, desc=f"Generating response from {model_path_1}")
+    response1 = generate_response(instruction, model_path_1, progress)
+    
+    progress(0.4, desc=f"Generating response from {model_path_2}")
+    response2 = generate_response(instruction, model_path_2, progress)
+    
+    # Update the response textboxes
+    progress(0.7, desc="Evaluating responses")
+    
+    # Use the judge_responses generator but collect all outputs
+    evaluation_results = ""
+    for result in judge_responses(instruction, response1, response2, judge_model_name, temperature, max_new_tokens):
+        evaluation_results = result
+    
+    progress(1.0, desc="Evaluation complete")
+    
+    return response1, response2, evaluation_results
 
+# Create Gradio interface
+with gr.Blocks(fill_height=True, css=css) as demo:
+    gr.Markdown(DESCRIPTION)
+    
+    with gr.Tabs():
+        # Auto-Generate Responses tab (now first)
+        with gr.TabItem("Auto-Generate Responses"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    # Model selection for judge
+                    auto_model_dropdown = gr.Dropdown(
+                        choices=list(MODEL_PATHS.keys()),
+                        value="JudgeLRM-7B",  # Default selection
+                        label="Select Judge Model"
+                    )
+                    
+                    auto_instruction = gr.Textbox(label="Instruction/Question", lines=3)
+                    
+                    # Model 1 selection
+                    with gr.Row():
+                        model_dropdown_1 = gr.Dropdown(
+                            choices=POPULAR_MODELS,
+                            value=POPULAR_MODELS[0],
+                            label="Select Model 1",
+                            scale=2
+                        )
+                        custom_model_1 = gr.Textbox(
+                            label="Or enter custom model path",
+                            placeholder="e.g., meta-llama/Llama-2-7b-chat-hf",
+                            scale=3
+                        )
+                    
+                    # Model 2 selection
+                    with gr.Row():
+                        model_dropdown_2 = gr.Dropdown(
+                            choices=POPULAR_MODELS,
+                            value=POPULAR_MODELS[1],
+                            label="Select Model 2",
+                            scale=2
+                        )
+                        custom_model_2 = gr.Textbox(
+                            label="Or enter custom model path",
+                            placeholder="e.g., mistralai/Mistral-7B-Instruct-v0.2",
+                            scale=3
+                        )
+                    
+                    with gr.Accordion("⚙️ Parameters", open=False):
+                        auto_temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.1, label="Judge Temperature")
+                        auto_max_tokens = gr.Slider(minimum=128, maximum=4096, step=1, value=2048, label="Judge Max Tokens")
+                    
+                    auto_submit_btn = gr.Button("Generate Responses and Evaluate")
+            
+            with gr.Row():
+                with gr.Column():
+                    auto_response1 = gr.Textbox(label="Response from Model 1", lines=10)
+                
+                with gr.Column():
+                    auto_response2 = gr.Textbox(label="Response from Model 2", lines=10)
+            
+            with gr.Row():
+                auto_output = gr.Textbox(label="Evaluation Results", lines=15)
+            
+            # Handle auto-generation and evaluation
+            auto_submit_btn.click(
+                fn=generate_and_judge,
+                inputs=[auto_instruction, model_dropdown_1, custom_model_1, model_dropdown_2, custom_model_2, auto_model_dropdown, auto_temperature, auto_max_tokens],
+                outputs=[auto_response1, auto_response2, auto_output]
+            )
+            # Examples for auto-generation
+            auto_examples = [
+                ["Write a short poem about artificial intelligence",
+                 "Qwen/Qwen2.5-7B-Instruct",
+                 "Qwen/Qwen2.5-7B-Instruct",
+                 "01-ai/Yi-6B-Chat",
+                 "01-ai/Yi-6B-Chat"],
+                ["我听说有些人有高血压却没有任何症状。这是真的吗？",
+                 "FreedomIntelligence/Apollo-7B",
+                 "FreedomIntelligence/Apollo-7B",
+                 "microsoft/phi-2",
+                 "microsoft/phi-2"]
+            ]
+            
+            gr.Examples(
+                examples=auto_examples,
+                inputs=[auto_instruction, model_dropdown_1, custom_model_1, model_dropdown_2, custom_model_2]
+            )
+        
+        # Manual Evaluation tab (now second)
+        with gr.TabItem("Manual Evaluation"):
+            with gr.Row():
+                with gr.Column():
+                    # Model selection
+                    model_dropdown = gr.Dropdown(
+                        choices=list(MODEL_PATHS.keys()),
+                        value="JudgeLRM-7B",  # Default selection
+                        label="Select Judge Model"
+                    )
+                    
+                    instruction = gr.Textbox(label="Instruction/Question", lines=3)
+                    response1 = gr.Textbox(label="Response 1", lines=8)
+                    response2 = gr.Textbox(label="Response 2", lines=8)
+                    
+                    with gr.Accordion("⚙️ Parameters", open=False):
+                        temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.1, label="Temperature")
+                        max_tokens = gr.Slider(minimum=128, maximum=4096, step=1, value=2048, label="Max Tokens")
+                    
+                    submit_btn = gr.Button("Evaluate Responses")
+                
+                with gr.Column():
+                    output = gr.Textbox(label="Evaluation Results", lines=20)
+            
+            # Handle evaluation
+            submit_btn.click(
+                fn=judge_responses,
+                inputs=[instruction, response1, response2, model_dropdown, temperature, max_tokens],
+                outputs=output
+            )
+            
+            # Examples
+            examples = [
+                ["Include important study notes and key points that someone should know about the given subject. \n history of the USA",
+                 "The history of the United States is one of the most influential and influential countries in the world. It has had a profound impact on the world and continues to shape the world today. The United States has had a long and storied history, from its founding in 1775 to the present day. It has had a tumultuous and tumultuous history, from the Revolutionary War to the Civil War and beyond. The United States has had a tumultuous and tumultuous history, from the Revolutionary War to the Civil War and beyond. The United States has had a long and storied history, from its founding in 1775 to the present day. It has had a profound impact on the world and continues to shape the world today.",
+                 "1. The United States of America was founded in 1776.\n2. The Declaration of Independence was signed in 1776.\n3. The Constitution of the United States of America was signed in 1787.\n4. The Civil War began in 1861.\n5. The Emancipation Proclamation was issued in 1863.\n6. The 13th Amendment was ratified in 1865.\n7. The 14th Amendment was ratified in 1868.\n8. The 15th Amendment was ratified in 1870.\n9. The 16th Amendment was ratified in 1913.\n10. The 17th Amendment was ratified in 1913.\n11. The 18th Amendment was ratified in 1919.\n12. The 19th Amendment was ratified in 1920.\n13. The 20th Amendment was ratified in 1933.\n14. The 21st Amendment was ratified in 1933.\n15. The 22nd Amendment was ratified in"]
+            ]
+            
+            gr.Examples(examples=examples, inputs=[instruction, response1, response2])
+    
+    gr.Markdown(LICENSE)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()