ltg/fluency-annotation
0
1from __future__ import annotations2 3import os4import gradio as gr5import json6import random7from datetime import datetime8from typing import Dict, List, Tuple9import hashlib10import itertools11from datasets import load_dataset, Dataset, DatasetDict12from huggingface_hub import HfApi, create_repo, repo_exists, Repository13from huggingface_hub import HfFolder14import shutil15import threading16import json17 18from collections.abc import Iterable19 20from gradio.themes.base import Base21from gradio.themes.utils import colors, fonts, sizes22 23HF_TOKEN = os.environ.get("HF_TOKEN")24os.environ['HF_AUTH'] = HF_TOKEN25HfApi(token=HF_TOKEN)26 27USER_IDS = set(json.loads(os.environ.get("USER_IDS")) + json.loads(os.environ.get("USER_IDS_2")))28 29 30class Soft(Base):31 def __init__(32 self,33 *,34 primary_hue: colors.Color | str = colors.indigo,35 secondary_hue: colors.Color | str = colors.indigo,36 neutral_hue: colors.Color | str = colors.gray,37 spacing_size: sizes.Size | str = sizes.spacing_md,38 radius_size: sizes.Size | str = sizes.radius_md,39 text_size: sizes.Size | str = sizes.text_md,40 font: fonts.Font | str | Iterable[fonts.Font | str] = (41 # fonts.LocalFont("Montserrat"),42 "ui-sans-serif",43 "system-ui",44 "sans-serif",45 ),46 font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (47 # fonts.LocalFont("IBM Plex Mono"),48 "ui-monospace",49 "Consolas",50 "monospace",51 ),52 ):53 super().__init__(54 primary_hue=primary_hue,55 secondary_hue=secondary_hue,56 neutral_hue=neutral_hue,57 spacing_size=spacing_size,58 radius_size=radius_size,59 text_size=text_size,60 font=font,61 font_mono=font_mono,62 )63 self.name = "soft"64 super().set(65 # Colors66 background_fill_primary="*neutral_50",67 slider_color="*primary_500",68 slider_color_dark="*primary_600",69 # Shadows70 shadow_drop="0 1px 4px 0 rgb(0 0 0 / 0.1)",71 shadow_drop_lg="0 2px 5px 0 rgb(0 0 0 / 0.2)",72 # Block Labels73 block_background_fill="white",74 block_label_padding="*spacing_sm *spacing_md",75 block_label_background_fill="*primary_100",76 block_label_background_fill_dark="*primary_600",77 block_label_radius="*radius_md",78 block_label_text_size="*text_md",79 block_label_text_weight="600",80 block_label_text_color="*primary_500",81 block_label_text_color_dark="white",82 block_title_radius="*block_label_radius",83 block_title_padding="*block_label_padding",84 block_title_background_fill="*block_label_background_fill",85 block_title_text_weight="600",86 block_title_text_color="*primary_500",87 block_title_text_color_dark="white",88 block_label_margin="*spacing_md",89 90 # Inputs91 input_background_fill="white",92 input_border_color="*neutral_100",93 input_shadow="*shadow_drop",94 input_shadow_focus="*shadow_drop_lg",95 checkbox_shadow="none",96 # Buttons97 shadow_spread="6px",98 button_primary_shadow="*shadow_drop_lg",99 button_primary_shadow_hover="*shadow_drop_lg",100 button_primary_shadow_active="*shadow_inset",101 button_secondary_shadow="*shadow_drop_lg",102 button_secondary_shadow_hover="*shadow_drop_lg",103 button_secondary_shadow_active="*shadow_inset",104 checkbox_label_shadow="*shadow_drop_lg",105 button_primary_background_fill="*primary_500",106 button_primary_background_fill_hover="*primary_400",107 button_primary_background_fill_hover_dark="*primary_500",108 button_primary_text_color="white",109 button_secondary_background_fill="white",110 button_secondary_background_fill_hover="*neutral_100",111 button_secondary_background_fill_hover_dark="*primary_500",112 button_secondary_text_color="*neutral_800",113 button_cancel_background_fill="*button_secondary_background_fill",114 button_cancel_background_fill_hover="*button_secondary_background_fill_hover",115 button_cancel_background_fill_hover_dark="*button_secondary_background_fill_hover",116 button_cancel_text_color="*button_secondary_text_color",117 checkbox_label_background_fill_selected="*primary_500",118 checkbox_label_background_fill_selected_dark="*primary_600",119 checkbox_border_width="1px",120 checkbox_border_color="*neutral_100",121 checkbox_border_color_dark="*neutral_600",122 checkbox_background_color_selected="*primary_600",123 checkbox_background_color_selected_dark="*primary_700",124 checkbox_border_color_focus="*primary_500",125 checkbox_border_color_focus_dark="*primary_600",126 checkbox_border_color_selected="*primary_600",127 checkbox_border_color_selected_dark="*primary_700",128 checkbox_label_text_color_selected="white",129 # Borders130 block_border_width="0px",131 panel_border_width="0px",132 )133 134 135guideline = open("guidelines.md").read().strip()136 137# Configuration for the output dataset138ANNOTATIONS_REPO = "ltg/fluency-annotations" # Change to your repo name139DATA_DIR = "annotation_data"140ANNOTATIONS_FILE = os.path.join(DATA_DIR, "train.jsonl")141 142# Model names for the three responses143MODEL_NAMES = ["mistral-Nemo", "translated-SFT", "on-policy-RL"]144 145# Create all pairwise comparisons146MODEL_PAIRS = list(itertools.combinations(MODEL_NAMES, 2))147 148# Initialize repository149def init_repository():150 """Initialize or clone the repository"""151 try:152 repo = Repository(153 local_dir=DATA_DIR, 154 clone_from=ANNOTATIONS_REPO, 155 use_auth_token=HF_TOKEN,156 repo_type="dataset"157 )158 repo.git_pull()159 return repo160 except Exception as e:161 print(f"Error initializing repository: {e}")162 # Create local directory if repo doesn't exist163 os.makedirs(DATA_DIR, exist_ok=True)164 return None165 166# Initialize on startup167annotation_repo = init_repository()168 169def load_existing_annotations():170 """Load existing annotations from the jsonl file"""171 annotations = {}172 173 if os.path.exists(ANNOTATIONS_FILE):174 try:175 with open(ANNOTATIONS_FILE, "r") as f:176 for line in f:177 if line.strip():178 ann = json.loads(line)179 user_id = ann.get("user_id")180 if user_id:181 if user_id not in annotations:182 annotations[user_id] = []183 annotations[user_id].append(ann)184 print(f"Loaded {sum(len(v) for v in annotations.values())} existing annotations")185 except Exception as e:186 print(f"Error loading annotations: {e}")187 188 return annotations189 190def save_annotation_to_file(annotation_data):191 """Save a single annotation to the jsonl file and push to hub"""192 global annotation_repo193 194 try:195 # Pull latest changes196 if annotation_repo:197 annotation_repo.git_pull()198 199 # Append to jsonl file200 with open(ANNOTATIONS_FILE, "a") as f:201 line = json.dumps(annotation_data, ensure_ascii=False)202 f.write(f"{line}\n")203 204 # Push to hub asynchronously205 if annotation_repo:206 annotation_repo.push_to_hub(blocking=False)207 208 except Exception as e:209 print(f"Error saving annotation: {e}")210 # Try to reinitialize repository211 try:212 shutil.rmtree(DATA_DIR)213 annotation_repo = init_repository()214 215 # Retry saving216 with open(ANNOTATIONS_FILE, "a") as f:217 line = json.dumps(annotation_data, ensure_ascii=False)218 f.write(f"{line}\n")219 220 if annotation_repo:221 annotation_repo.push_to_hub(blocking=False)222 except Exception as e2:223 print(f"Failed to save annotation after retry: {e2}")224 225def load_dataset_samples():226 """Load and prepare dataset samples with pairwise comparisons"""227 try:228 # Load the private dataset (requires authentication)229 dataset = load_dataset("ltg/fluency-generations", split="train", token=HF_TOKEN)230 231 # Transform dataset into pairwise comparison format232 pairwise_samples = []233 234 for item in dataset:235 sample_id = item["sample_id"]236 prompt = item["prompt"]237 responses = item["responses"]238 239 # Create pairwise comparisons for this sample240 for model_a, model_b in MODEL_PAIRS:241 pairwise_samples.append({242 "id": f"{sample_id}_{model_a}_vs_{model_b}",243 "original_id": sample_id,244 "prompt": prompt,245 "response_a": responses[model_a],246 "response_b": responses[model_b],247 "model_a": model_a,248 "model_b": model_b,249 "dataset": "NTNU"250 })251 252 extra_dataset = load_dataset("ltg/fluency-generations", split="test", token=HF_TOKEN)253 extra_pairwise_samples = []254 for i, item in enumerate(extra_dataset):255 sample_id = item["sample_id"]256 prompt = item["prompt"]257 responses = item["responses"]258 model_a, model_b = MODEL_PAIRS[i % len(MODEL_PAIRS)]259 model_a, model_b = (model_a, model_b) if i % 2 == 0 else (model_b, model_a)260 extra_pairwise_samples.append({261 "id": f"{sample_id}_{model_a}_vs_{model_b}",262 "original_id": sample_id,263 "prompt": prompt,264 "response_a": responses[model_a],265 "response_b": responses[model_b],266 "model_a": model_a,267 "model_b": model_b,268 "dataset": "training_examples"269 })270 271 return pairwise_samples, extra_pairwise_samples272 273 except Exception as e:274 print(f"Error loading dataset: {e}")275 print("Using dummy data for testing...")276 # Fallback to dummy data for testing277 return [278 {279 "id": "dummy_001_modelA_vs_modelB",280 "original_id": "dummy_001",281 "prompt": "Test prompt for development",282 "response_a": "This is response A for testing.",283 "response_b": "This is response B for testing.",284 "model_a": "modelA",285 "model_b": "modelB",286 "dataset": "test"287 }288 ], []289 290 291def swap_sample(sample):292 return {293 "id": str(sample["original_id"]) + '_' + sample["model_b"] + '_vs_' + sample["model_a"],294 "original_id": sample["original_id"],295 "prompt": sample["prompt"],296 "response_a": sample["response_b"],297 "response_b": sample["response_a"],298 "model_a": sample["model_b"],299 "model_b": sample["model_a"],300 "dataset": sample["dataset"]301 }302 303# Load dataset on startup304DATASET_SAMPLES, EXTRA_DATASET_SAMPLES = load_dataset_samples()305 306class AnnotationManager:307 def __init__(self):308 # Load existing annotations from file309 self.annotations = load_existing_annotations()310 self.user_states = {}311 312 # Rebuild user states from loaded annotations313 for user_id, user_annotations in self.annotations.items():314 annotated_ids = [ann["sample_id"] for ann in user_annotations]315 self.user_states[user_id] = {316 "current_index": 0,317 "annotations": annotated_ids318 }319 320 def get_user_seed(self, user_id: str) -> int:321 """Generate consistent seed for user"""322 return int(hashlib.md5(user_id.encode()).hexdigest(), 16)323 324 def get_user_samples(self, user_id: str) -> List[Dict]:325 """Get shuffled samples for user based on their ID"""326 seed = self.get_user_seed(user_id)327 samples = DATASET_SAMPLES.copy()328 random.Random(seed).shuffle(samples)329 samples = [330 sample if random.Random(seed + i).randint(0, 1) == 0 else swap_sample(sample)331 for i, sample in enumerate(samples)332 ]333 samples = EXTRA_DATASET_SAMPLES.copy() + samples334 return samples335 336 def get_next_sample(self, user_id: str) -> Tuple[Dict, int, int]:337 """Get next unannotated sample for user"""338 if user_id not in self.user_states:339 # Check if user has existing annotations340 if user_id in self.annotations:341 annotated_ids = [ann["sample_id"] for ann in self.annotations[user_id]]342 self.user_states[user_id] = {343 "current_index": 0,344 "annotations": annotated_ids345 }346 else:347 self.user_states[user_id] = {348 "current_index": 0,349 "annotations": []350 }351 352 samples = self.get_user_samples(user_id)353 state = self.user_states[user_id]354 355 # Count total annotations for this user356 total_annotated = len(state["annotations"])357 358 # Find next unannotated sample359 for idx, sample in enumerate(samples):360 if not self.is_annotated(user_id, sample["id"]):361 return sample, total_annotated + 1, len(samples)362 363 # All samples annotated364 return None, len(samples), len(samples)365 366 def is_annotated(self, user_id: str, sample_id: str) -> bool:367 """Check if user has annotated this sample"""368 if user_id not in self.annotations:369 return False370 return any(ann["sample_id"] == sample_id for ann in self.annotations[user_id])371 372 def save_annotation(self, user_id: str, sample_id: str, choice: str, 373 model_a: str = None, model_b: str = None, 374 original_id: str = None, dataset_name: str = None):375 """Save user's annotation and persist to file"""376 if user_id not in self.annotations:377 self.annotations[user_id] = []378 379 annotation = {380 "user_id": user_id,381 "sample_id": sample_id,382 "original_sample_id": original_id,383 "dataset": dataset_name,384 "model_a": model_a,385 "model_b": model_b,386 "choice": choice,387 "timestamp": datetime.now().isoformat()388 }389 390 # Save to memory391 self.annotations[user_id].append(annotation)392 393 # Update user state394 if user_id in self.user_states:395 self.user_states[user_id]["annotations"].append(sample_id)396 else:397 self.user_states[user_id] = {398 "current_index": 0,399 "annotations": [sample_id]400 }401 402 # Save to file asynchronously403 threading.Thread(404 target=save_annotation_to_file, 405 args=(annotation,)406 ).start()407 408 print(f"Saved annotation: {annotation}")409 410 def get_user_progress(self, user_id: str) -> Dict:411 """Get user's annotation progress"""412 if user_id not in self.annotations:413 return {"completed": 0, "total": len(DATASET_SAMPLES)}414 415 completed = len(self.annotations[user_id])416 return {"completed": completed, "total": len(DATASET_SAMPLES)}417 418 419# Initialize manager420manager = AnnotationManager()421 422def login(user_id: str) -> Tuple:423 """Handle user login"""424 if not user_id or user_id.strip() == "" or user_id.strip() not in USER_IDS:425 return (426 gr.update(visible=True), # login_interface427 gr.update(visible=False), # annotation_interface428 "", # user_state429 gr.update(value="Please enter a valid ID"), # login_status430 gr.update(), # prompt431 gr.update(), # response_a432 gr.update(), # response_b433 gr.update() # progress434 )435 436 user_id = user_id.strip()437 sample, current, total = manager.get_next_sample(user_id)438 439 if sample is None:440 return (441 gr.update(visible=True), # login_interface442 gr.update(visible=False), # annotation_interface443 user_id, # user_state444 gr.update(value=f"All {total} samples completed for user: {user_id}! 🎉"), # login_status445 gr.update(), # prompt446 gr.update(), # response_a447 gr.update(), # response_b448 gr.update() # progress449 )450 451 # Show which models are being compared452 model_info = f" | Comparing: {sample.get('model_a', 'A')} vs {sample.get('model_b', 'B')}"453 454 return (455 gr.update(visible=False), # login_interface456 gr.update(visible=True), # annotation_interface457 user_id, # user_state458 gr.update(value=""), # login_status459 gr.update(value=sample["prompt"]), # prompt460 gr.update(value=sample["response_a"]), # response_a461 gr.update(value=sample["response_b"]), # response_b462 gr.update(value=f"Progress: {current}/{total}") # progress463 )464 465def annotate(choice: str, user_id: str) -> Tuple:466 """Handle annotation submission"""467 if not user_id:468 return (469 gr.update(), # prompt470 gr.update(), # response_a471 gr.update(), # response_b472 gr.update(), # progress473 gr.update(value="Error: No user logged in", visible=True) # status474 )475 476 # Get current sample to save annotation477 sample, _, _ = manager.get_next_sample(user_id)478 if sample:479 # Map button choice to annotation value480 choice_map = {481 "a_better": "A is more fluent",482 "b_better": "B is more fluent",483 "equal": "Equally fluent"484 }485 # Save with all metadata486 manager.save_annotation(487 user_id=user_id,488 sample_id=sample["id"],489 choice=choice_map[choice],490 model_a=sample.get("model_a"),491 model_b=sample.get("model_b"),492 original_id=sample.get("original_id"),493 dataset_name=sample.get("dataset")494 )495 496 # Get next sample497 next_sample, current, total = manager.get_next_sample(user_id)498 499 if next_sample is None:500 return (501 gr.update(value="All samples completed! Thank you for your annotations."), # prompt502 gr.update(value=""), # response_a503 gr.update(value=""), # response_b504 gr.update(value=f"Progress: {total}/{total} - Complete!"), # progress505 gr.update(value="All annotations complete!", visible=True) # status506 )507 508 # Show which models are being compared509 model_info = f" | Comparing: {next_sample.get('model_a', 'A')} vs {next_sample.get('model_b', 'B')}"510 511 return (512 gr.update(value=next_sample["prompt"]), # prompt513 gr.update(value=next_sample["response_a"]), # response_a514 gr.update(value=next_sample["response_b"]), # response_b515 gr.update(value=f"Progress: {current}/{total}"), # progress516 gr.update(value="Annotation saved!", visible=True) # status517 )518 519def logout() -> Tuple:520 """Handle user logout"""521 return (522 gr.update(visible=True), # login_interface523 gr.update(visible=False), # annotation_interface524 "", # user_state525 gr.update(value=""), # login_status526 gr.update(value=""), # prompt527 gr.update(value=""), # response_a528 gr.update(value=""), # response_b529 gr.update(value="") # progress530 )531 532# Create Gradio interface533custom_css = """534 #login-group {535 background-color: white !important;536 }537 #login-group > * {538 background-color: white !important;539 }540 #login-group .gr-group {541 background-color: white !important;542 }543 #login-group .gr-form {544 background-color: white !important;545 }546 .light-shadow {547 box-shadow: 0 1px 4px 0 rgb(0 0 0 / 0.1) !important;548 }549 /* Target the textbox container */550 .no-style-textbox {551 border: none !important;552 box-shadow: none !important;553 }554 555 /* Target both input and textarea elements */556 .no-style-textbox input,557 .no-style-textbox textarea {558 border: none !important;559 box-shadow: none !important;560 padding: 0 !important;561 outline: none !important;562 }563 564 /* Target the Gradio textbox wrapper */565 .no-style-textbox .gr-textbox {566 border: none !important;567 box-shadow: none !important;568 }569 570 /* Target focus states */571 .no-style-textbox input:focus,572 .no-style-textbox textarea:focus {573 border: none !important;574 box-shadow: none !important;575 outline: none !important;576 }577 578 /* Additional targeting for stubborn Gradio elements */579 .no-style-textbox .gr-form,580 .no-style-textbox .gr-input {581 border: none !important;582 box-shadow: none !important;583 }584"""585 586# Create Gradio interface587with gr.Blocks(theme=Soft(font=[gr.themes.GoogleFont("Source Sans Pro"), "Arial"]), title="Dataset Annotation Tool", css=custom_css) as app:588 gr.Markdown("# Norwegian Fluency Annotation")589 with gr.Accordion("Click here to see the full annotation guidelines:", open=False, elem_classes="light-shadow"):590 gr.Markdown(guideline, padding=True)591 592 user_state = gr.State("")593 594 # Login Interface595 with gr.Column(visible=True) as login_interface:596 with gr.Column(variant="panel", elem_id="login-group", elem_classes="light-shadow"):597 gr.Markdown("## Log in", padding=True)598 user_id_input = gr.Textbox(599 label="Enter your unique annotator ID to begin",600 placeholder="Annotator ID"601 )602 with gr.Row():603 login_btn = gr.Button("Login", variant="primary", scale=0.2, min_width=100)604 gr.HTML("")605 login_status = gr.Markdown("", padding=True)606 607 # Annotation Interface608 with gr.Column(visible=False, elem_id="annotation-group") as annotation_interface:609 progress_label = gr.Markdown("")610 611 # Row 1: Prompt612 with gr.Row(elem_classes="light-shadow"):613 prompt_display = gr.Textbox(614 label="Prompt",615 interactive=False,616 lines=1,617 elem_classes="no-style-textbox",618 autoscroll=False619 )620 621 # Row 2: Responses622 with gr.Row(elem_classes="light-shadow"):623 response_a_display = gr.Textbox(624 label="Response A",625 interactive=False,626 lines=1,627 scale=1,628 elem_classes="no-style-textbox",629 autoscroll=False,630 max_lines=100631 )632 response_b_display = gr.Textbox(633 label="Response B",634 interactive=False,635 lines=1,636 scale=1,637 elem_classes="no-style-textbox",638 autoscroll=False,639 max_lines=100640 )641 642 # Row 3: Buttons643 with gr.Row():644 btn_a = gr.Button("A is more fluent", variant="primary")645 btn_equal = gr.Button("Equally fluent", variant="primary")646 btn_b = gr.Button("B is more fluent", variant="primary")647 648 status_message = gr.Markdown("", visible=False)649 650 with gr.Row(visible=False):651 logout_btn = gr.Button("Logout", variant="stop", size="sm")652 653 # Event handlers654 login_btn.click(655 fn=login,656 inputs=[user_id_input],657 outputs=[658 login_interface, 659 annotation_interface, 660 user_state, 661 login_status,662 prompt_display,663 response_a_display,664 response_b_display,665 progress_label666 ]667 )668 669 user_id_input.submit(670 fn=login,671 inputs=[user_id_input],672 outputs=[673 login_interface, 674 annotation_interface, 675 user_state, 676 login_status,677 prompt_display,678 response_a_display,679 response_b_display,680 progress_label681 ]682 )683 684 btn_a.click(685 fn=lambda user_id: annotate("a_better", user_id),686 inputs=[user_state],687 outputs=[688 prompt_display,689 response_a_display,690 response_b_display,691 progress_label,692 status_message693 ]694 )695 696 btn_b.click(697 fn=lambda user_id: annotate("b_better", user_id),698 inputs=[user_state],699 outputs=[700 prompt_display,701 response_a_display,702 response_b_display,703 progress_label,704 status_message705 ]706 )707 708 btn_equal.click(709 fn=lambda user_id: annotate("equal", user_id),710 inputs=[user_state],711 outputs=[712 prompt_display,713 response_a_display,714 response_b_display,715 progress_label,716 status_message717 ]718 )719 720 logout_btn.click(721 fn=logout,722 inputs=[],723 outputs=[724 login_interface,725 annotation_interface,726 user_state,727 login_status,728 prompt_display,729 response_a_display,730 response_b_display,731 progress_label732 ]733 )734 735if __name__ == "__main__":736 app.launch()