app.py

import gradio as gr
import random
import json
import os
import time
from datetime import datetime
from datasets import load_dataset, Dataset
from huggingface_hub import HfApi, create_repo, dataset_info
import pandas as pd

from dotenv import load_dotenv

load_dotenv()

# Configuration
DATASET_NAME = "aaronsnoswell/alignment-annotation-pairwise"
HF_TOKEN = os.getenv("HF_TOKEN")

ANNOTATION_GUIDELINES = """## Guidelines for annotators:
      
In making your choice, consider the following aspects of responses:

* **Honesty:**
The assistant should be honest about whether it knows the answer and express its uncertainty explicitly. The Assistant should be confident on questions it knows well and modest on those it is unfamiliar with.
The assistant should use weakeners such as "I guess", "I suppose", "probably", and "perhaps" to express uncertainty, and assistants should answer "I don't know" if necessary.

* **Truthfulness:**
The assistant should answer truthfully and be faithful to factual knowledge as well as given contexts, never making up any new facts that aren't true or cannot be grounded in the instruction.

* **Helpfulness:**
The assistant should provide users with accurate, relevant, and up-to-date information, ensuring that the content is positive, interesting, engaging, educational, and helpful.
"""

# Initialize HF API
api = HfApi()

# Load the source dataset
print("Loading UltraFeedback dataset...")
ds = load_dataset("openbmb/UltraFeedback")
train_data = ds['train']
print(f"Dataset loaded with {len(train_data)} examples")

def initialize_dataset():
    """Initialize the annotations dataset if it doesn't exist"""
    try:
        # Check if dataset exists
        dataset_info(DATASET_NAME, token=HF_TOKEN)
        print(f"Dataset {DATASET_NAME} already exists")
    except:
        # Create new dataset
        try:
            create_repo(
                repo_id=DATASET_NAME,
                repo_type="dataset",
                token=HF_TOKEN,
                exist_ok=True
            )
            
            # Create initial empty dataset
            initial_data = {
                "timestamp": [],
                "source_idx": [],
                "instruction": [],
                "completion_1": [],
                "completion_2": [],
                "preference": [],  # "left" or "right" - we don't save on "skip"
                "source_dataset": []
            }
            
            initial_df = pd.DataFrame(initial_data)
            initial_dataset = Dataset.from_pandas(initial_df)
            initial_dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
            print(f"Created new dataset: {DATASET_NAME}")
            
        except Exception as e:
            print(f"Error creating dataset: {e}")

def save_annotation(source_idx, instruction, completion_1, completion_2, preference):
    """Save an annotation to the HuggingFace dataset"""
    if not HF_TOKEN:
        print("No HF_TOKEN found - annotation not saved")
        return False
    
    try:
        # Prepare the annotation data
        annotation = {
            "timestamp": [datetime.now().isoformat()],
            "source_idx": [source_idx],
            "instruction": [instruction],
            "completion_1": [completion_1],
            "completion_2": [completion_2],
            "preference": [preference],
            "source_dataset": ["openbmb/UltraFeedback"]
        }
        
        # Create dataset from the annotation
        new_data = Dataset.from_dict(annotation)
        
        # Load existing dataset and concatenate
        try:
            existing_dataset = load_dataset(DATASET_NAME, token=HF_TOKEN, split="train")
            combined_dataset = Dataset.from_dict({
                **existing_dataset.to_dict(),
                **{k: existing_dataset[k] + v for k, v in annotation.items()}
            })
        except:
            # If dataset doesn't exist or is empty, use the new data
            combined_dataset = new_data
        
        # Push to hub
        combined_dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
        print(f"Saved annotation: {preference} preference for example {source_idx}")
        return True
        
    except Exception as e:
        print(f"Error saving annotation: {e}")
        return False

def get_random_example():
    """Get a random example from the dataset and format it for display"""
    idx = random.randint(0, len(train_data) - 1)
    dat = train_data[idx]
    
    source = dat['source']
    instruction = dat['instruction']
    models = dat['models']
    completions = dat['completions']
    
    # Get first two completions
    completion_1 = completions[0]['response']
    completion_2 = completions[1]['response']
    model_1 = "Completion A"
    model_2 = "Completion B"
    
    # Format prompt display
    prompt_display = f"## Prompt:\n\n{instruction}\n\n---"
    
    # Format completion displays
    completion_1_display = f"## {model_1}\n\n{completion_1}"
    completion_2_display = f"## {model_2}\n\n{completion_2}"

    print("Randomly loaded example: ", idx)
    
    return prompt_display, completion_1_display, completion_2_display, idx, instruction, completion_1, completion_2

def format_stats_display(judgment_times, num_judgments, num_skips):
    """Format the statistics display"""
    if num_judgments == 0:
        return "📊 **Session Statistics:** No judgments made yet."
    
    avg_time = sum(judgment_times) / len(judgment_times)
    
    stats = f"""📊 **Session Statistics:** {num_judgments} judgements made, {num_skips} items skipped. Average time per judgement {avg_time:.1f} seconds)."""
    
    return stats

def load_first_example():
    """Load the first example and start the annotation interface"""
    prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2 = get_random_example()
    start_time = time.time()
    
    return (
        prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2,
        start_time,
        gr.update(visible=False),  # Hide load button
        gr.update(visible=True),   # Show prompt
        gr.update(visible=True),   # Show completion row
        gr.update(visible=True),   # Show action buttons
    )

def handle_left_better(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2, 
                      start_time, judgment_times, num_judgments, num_skips):
    """Handle when user selects left completion as better"""
    print(f"User selected LEFT completion as better for example {current_idx}")
    
    # Calculate time taken for this judgment
    end_time = time.time()
    time_taken = end_time - start_time
    judgment_times.append(time_taken)
    num_judgments += 1
    
    print(f"Time taken for judgment: {time_taken:.1f} seconds")
    
    # Save the annotation
    success = save_annotation(current_idx, instruction, completion_1, completion_2, "left")
    
    # Get new random example
    new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
    
    # Reset timer for new example
    new_start_time = time.time()
    
    # Update stats display
    stats_display = format_stats_display(judgment_times, num_judgments, num_skips)

    message = "✅ Annotation saved! Left completion selected as better." if success else "✅ Left completion selected (save failed - check console)"
    gr.Info(message)
    
    return (
        new_prompt, 
        new_comp_1, 
        new_comp_2, 
        new_idx,
        new_instruction,
        new_completion_1,
        new_completion_2,
        new_start_time,
        judgment_times,
        num_judgments,
        num_skips,
        stats_display
    )

def handle_right_better(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2,
                       start_time, judgment_times, num_judgments, num_skips):
    """Handle when user selects right completion as better"""
    print(f"User selected RIGHT completion as better for example {current_idx}")
    
    # Calculate time taken for this judgment
    end_time = time.time()
    time_taken = end_time - start_time
    judgment_times.append(time_taken)
    num_judgments += 1
    
    print(f"Time taken for judgment: {time_taken:.1f} seconds")
    
    # Save the annotation
    success = save_annotation(current_idx, instruction, completion_1, completion_2, "right")
    
    # Get new random example
    new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
    
    # Reset timer for new example
    new_start_time = time.time()
    
    # Update stats display
    stats_display = format_stats_display(judgment_times, num_judgments, num_skips)

    message = "✅ Annotation saved! Right completion selected as better." if success else "✅ Right completion selected (save failed - check console)"
    gr.Info(message)
    
    return (
        new_prompt, 
        new_comp_1, 
        new_comp_2, 
        new_idx,
        new_instruction,
        new_completion_1,
        new_completion_2,
        new_start_time,
        judgment_times,
        num_judgments,
        num_skips,
        stats_display
    )

def handle_skip(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2,
               start_time, judgment_times, num_judgments, num_skips):
    """Handle when user skips the current example"""
    print(f"User skipped example {current_idx}")
    
    # Increment skip counter (don't track time for skips)
    num_skips += 1
    
    # Get new random example
    new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
    
    # Reset timer for new example
    new_start_time = time.time()
    
    # Update stats display
    stats_display = format_stats_display(judgment_times, num_judgments, num_skips)
    
    gr.Info("⏭️ Skipped example (not saved).")

    return (
        new_prompt, 
        new_comp_1, 
        new_comp_2, 
        new_idx,
        new_instruction,
        new_completion_1,
        new_completion_2,
        new_start_time,
        judgment_times,
        num_judgments,
        num_skips,
        stats_display
    )

# Initialize dataset on startup
if HF_TOKEN:
    initialize_dataset()
else:
    print("Warning: No HF_TOKEN found. Annotations will not be saved.")

def load_first_example():
    """Load the first example and start the annotation interface"""
    prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2 = get_random_example()
    start_time = time.time()
    
    return (
        prompt, comp_1, comp_2, idx, instruction, completion_1, completion_2,
        start_time,
        gr.update(visible=False),  # Hide load button
        gr.update(visible=True),   # Show prompt
        gr.update(visible=True),   # Show completion row
        gr.update(visible=True),   # Show action buttons
    )

# Create Gradio interface

with gr.Blocks(title="AI Alignment: Binary Preference Annotation", css=".square-button { height: 80px !important; }") as demo:
    gr.Markdown(f"""
# 🎯 AI Alignment: Binary Preference Annotation

You'll see a prompt and two AI completions. Select which completion you think is better, or skip if you're unsure.

This simulates the data annotation process used in RLHF (Reinforcement Learning from Human Feedback) training.

{ANNOTATION_GUIDELINES}
---
""")
    
    # State to track current example and its components
    current_idx = gr.State(0)
    current_instruction = gr.State("")
    current_completion_1 = gr.State("")
    current_completion_2 = gr.State("")
    
    # State to track timing and statistics
    start_time = gr.State(0.0)  # When current example was loaded
    judgment_times = gr.State([])  # List of times taken for each judgment
    num_judgments = gr.State(0)  # Number of judgments made
    num_skips = gr.State(0)  # Number of examples skipped
    
    # Load first example button (shown initially)
    load_first_btn = gr.Button("🚀 Load First Example", variant="primary", size="lg", elem_classes="square-button")
    
    # Display prompt (hidden initially)
    prompt_display = gr.Markdown("", label="Prompt", visible=False)
    
    # Display completions side by side (hidden initially)
    completion_row = gr.Row(visible=False)
    with completion_row:
        with gr.Column():
            completion_1_display = gr.Markdown("", label="Completion A (Left)")
        with gr.Column():
            completion_2_display = gr.Markdown("", label="Completion B (Right)")
    
    # Action buttons (hidden initially)
    action_buttons = gr.Row(visible=False)
    with action_buttons:
        left_better_btn = gr.Button("👈 Left is Better", variant="primary", size="lg")
        skip_btn = gr.Button("⏭️ Skip This Example", variant="secondary", size="lg")
        right_better_btn = gr.Button("👉 Right is Better", variant="primary", size="lg")
    
    # Add info about dataset saving
    status_msg = "**Status:** ❌ Not connected (annotations will not be saved)."
    if HF_TOKEN:
        status_msg = f"**Status:** ✅ Connected. Annotations are being saved to [{DATASET_NAME}](https://huggingface.co/datasets/{DATASET_NAME})"
    gr.Markdown(status_msg)
    
    # Statistics display
    stats_display = gr.Markdown("📊 **Session Statistics:** No judgments made yet.", label="Performance Stats")
    
    # Wire up the load first example button
    load_first_btn.click(
        load_first_example,
        inputs=[],
        outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                 start_time, load_first_btn, prompt_display, completion_row, action_buttons]
    )
    
    # Wire up the action buttons
    left_better_btn.click(
        handle_left_better,
        inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                start_time, judgment_times, num_judgments, num_skips],
        outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                 start_time, judgment_times, num_judgments, num_skips, stats_display]
    )
    
    right_better_btn.click(
        handle_right_better,
        inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                start_time, judgment_times, num_judgments, num_skips],
        outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                 start_time, judgment_times, num_judgments, num_skips, stats_display]
    )
    
    skip_btn.click(
        handle_skip,
        inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                start_time, judgment_times, num_judgments, num_skips],
        outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2,
                 start_time, judgment_times, num_judgments, num_skips, stats_display]
    )

if __name__ == "__main__":
    demo.launch()