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()