import gradio as gr import json import os import csv import pandas as pd from datasets import load_dataset, Dataset # Configuration DATASET_NAME = "anna-tch/generation-results" PROGRESS_FILE = "progress.csv" HF_TOKEN = os.getenv("HF_TOKEN") DEFAULT_ANNOTATION_COL = "manual_annotation" # Renamed for clarity def login(username, password, state): if password == os.getenv("LOGIN_PASSWORD"): # Store username and annotation column in state state["username"] = username new_annotation_col = f"{DEFAULT_ANNOTATION_COL}.{username}" state["annotation_col"] = new_annotation_col # Load data AFTER login with the correct annotation column data = load_data(state["annotation_col"]) state.update(data) # Merge data into the state return ( gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), state ) else: return ( gr.update(visible=True), gr.update(visible=False), gr.update(value="Incorrect password", visible=True), state ) def load_data(annotation_col): # Accept annotation_col as a parameter dataset = load_dataset(DATASET_NAME, token=HF_TOKEN)['train'] df = dataset.to_pandas().set_index('comment_id', drop=False) # Create user-specific annotation column if missing if annotation_col not in df.columns: df[annotation_col] = "" # Identify adapter columns dynamically (exclude annotation cols) adapter_columns = [ col for col in df.columns if col not in ['comment_id','inputs', 'mappings',annotation_col] and not col.startswith(DEFAULT_ANNOTATION_COL) ] # Track annotated/unannotated IDs annotated_ids = df[df[annotation_col] != ""].index.tolist() unannotated_ids = df[df[annotation_col] == ""].index.tolist() return { "annotation_col": annotation_col, "original_dataset": dataset, "dataset_df": df, "adapter_columns": adapter_columns, "annotated_ids": annotated_ids, "unannotated_ids": unannotated_ids, "current_view": "unannotated", "current_id_index": 0, "annotations": {} } def save_progress(annotation_data): with open(PROGRESS_FILE, 'a', newline='') as f: writer = csv.writer(f) row = [annotation_data['comment_id']] for adapter in annotation_data['scores']: row.extend([adapter, annotation_data['scores'][adapter]['grammar'], annotation_data['scores'][adapter]['coherence'], annotation_data['scores'][adapter]['hallucination'], annotation_data['scores'][adapter]['preferred_text'] ]) writer.writerow(row) def update_dataset(state): updated_dataset = Dataset.from_pandas( state["dataset_df"].reset_index(drop=True), ) updated_dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN) def get_current_comment(state): if state["current_view"] == "annotated": ids = state["annotated_ids"] else: ids = state["unannotated_ids"] if not ids: return None, None comment_id = ids[state["current_id_index"]] return state["dataset_df"].loc[comment_id], comment_id def update_display(state): example, comment_id = get_current_comment(state) if not comment_id: return [gr.update(value="No examples!")] + [gr.update(visible=False)] * len(state["adapter_columns"]) * 5 + [gr.update(value="Complete!"), state] # Use annotation column from state annotation_col = state["annotation_col"] existing_annotation = {} if example[annotation_col]: existing_annotation = json.loads(example[annotation_col]) # Retrieve the original inputs input_text = example.get("inputs", "No input available") updates = [gr.update(value=input_text)] # First element is the inputs textbox update for adapter in state["adapter_columns"]: text = example.get(adapter, "") scores = existing_annotation.get(adapter, {}) if existing_annotation else {} updates.extend([ gr.update(value=text, visible=True), gr.update(value=scores.get("grammar"), visible=True), gr.update(value=scores.get("coherence"), visible=True), gr.update(value=scores.get("hallucination"), visible=True), gr.update(value=scores.get("preferred_text"), visible=True) ]) return updates + [ gr.update(value=f"Example {state['current_id_index'] + 1} of {len(state[state['current_view'] + '_ids'])} ({state['current_view']})"), state ] def navigate(direction, state): ids = state[state["current_view"] + "_ids"] max_index = len(ids) - 1 if direction == "next" and state["current_id_index"] < max_index: state["current_id_index"] += 1 elif direction == "prev" and state["current_id_index"] > 0: state["current_id_index"] -= 1 return update_display(state) def switch_view(target_view, state): state["current_view"] = target_view state["current_id_index"] = 0 return update_display(state) def submit(*args): state = args[-1] annotation_col = state["annotation_col"] comment_id = state[state["current_view"] + "_ids"][state["current_id_index"]] # Each adapter has 9 components: text, grammar_title, grammar, coherence_title, coherence, separator, hallucination_title, hallucination, separator components_per_adapter = 9 adapter_scores = {} gramar_scores = args[:len(state["adapter_columns"])] coherence_scores = args[len(state["adapter_columns"]):len(state["adapter_columns"])*2] hallucination_scores = args[len(state["adapter_columns"])*2:len(state["adapter_columns"])*3] preferred_texts = args[len(state["adapter_columns"])*3:] for i in range(len(state["adapter_columns"])): #base_idx print(i) try: grammar = gramar_scores[i] print(state["adapter_columns"][i]) print("Grammar: ", grammar) coherence = coherence_scores[i] print("Coherence: ", coherence) hallucination = hallucination_scores[i] print("Hallucination: ", hallucination) preferred_text = preferred_texts[i] print("Preferred Text: ", preferred_text) adapter_scores[state["adapter_columns"][i]] = { "grammar": int(grammar) if grammar else None, "coherence": int(coherence) if coherence else None, "hallucination": int(hallucination) if hallucination else None, "preferred_text": preferred_text } except IndexError: print("Error: Not all scores provided") annotation = { "comment_id": comment_id, "scores": adapter_scores } # Update state state["annotations"][comment_id] = annotation state["dataset_df"].at[comment_id, annotation_col] = json.dumps(adapter_scores) #state["dataset_df"].at[comment_id, ANNOTATION_COL] = json.dumps(adapter_scores) # Update lists if new annotation if comment_id in state["unannotated_ids"]: state["unannotated_ids"].remove(comment_id) state["annotated_ids"].append(comment_id) save_progress(annotation) update_dataset(state) return navigate("next", state) with gr.Blocks() as app: # Initialize state as an empty dict state = gr.State(load_data(DEFAULT_ANNOTATION_COL)) # Login interface with gr.Row(visible=True, elem_id="login_row") as login_row: with gr.Column(): gr.Markdown("## Login") username_input = gr.Textbox(label="Username") password_input = gr.Textbox(label="Password", type="password") login_button = gr.Button("Login") login_error = gr.Markdown(visible=False) # Main interface with gr.Column(visible=False, elem_id="main_interface") as main_interface: gr.Markdown("## Text Annotation Tool") print("===>", state.value["annotation_col"]) with gr.Row(): # Sidebar Column with gr.Column(scale=1.5, min_width=300, elem_classes="sidebar"): gr.Markdown("### Navigation") with gr.Column(elem_classes="nav-buttons"): annotated_btn = gr.Button("Annotated", variant="primary") unannotated_btn = gr.Button("Unannotated", variant="primary") gr.Markdown("---") counter = gr.Markdown() # Main Content Column with gr.Column(scale=5, min_width=600, variant="compact"): # Add a textbox for the input at the top with gr.Accordion('Original Input', open=False): input_display = gr.Markdown() #input_display = gr.Textbox(interactive=False) gr.Markdown("---") adapter_components = [] for adapter in state.value["adapter_columns"]: with gr.Row(): adapter_components.extend([ gr.Textbox(label=f"{adapter}", visible=False, interactive=False, scale=2, lines=9), gr.Radio(choices=[1, 2, 3, 4, 5], label=f"Grammar", visible=False, scale=0.3, min_width=100), gr.Radio(choices=[1, 2, 3, 4, 5], label=f"Coherence", visible=False, scale=0.3, min_width=100), gr.Radio(choices=[1, 2, 3, 4, 5], label=f"Hallucination", visible=False, scale=0.3, min_width=100), gr.Checkbox(label="Preferred", visible=False, scale=0.2) ]) gr.Markdown("---") with gr.Row(): prev_btn = gr.Button("Previous", variant="secondary") next_btn = gr.Button("Next", variant="secondary") submit_btn = gr.Button("Submit", variant="primary") app.css = """ .sidebar { background: #f8f9fa; padding: 10px; border-right: 1px solid #dee2e6; height: 100vh; } .nav-buttons { background: #f0f0f0; padding: 20px; border-radius: 10px; } #login_row { padding: 20px; } """ # Login handler login_button.click( login, inputs=[username_input, password_input, state], outputs=[login_row, main_interface, login_error, state] ) navigation_inputs = [state] navigation_outputs = [input_display] + adapter_components + [counter, state] annotated_btn.click( lambda s: switch_view("annotated", s), inputs=[state], outputs=navigation_outputs ) unannotated_btn.click( lambda s: switch_view("unannotated", s), inputs=[state], outputs=navigation_outputs ) prev_btn.click( lambda s: navigate("prev", s), inputs=navigation_inputs, outputs=navigation_outputs ) next_btn.click( lambda s: navigate("next", s), inputs=navigation_inputs, outputs=navigation_outputs ) submit_btn.click( submit, inputs=adapter_components[1::5] + adapter_components[2::5] + adapter_components[3::5] + adapter_components[4::5] + [state], outputs=navigation_outputs ) app.load( lambda: update_display(state.value), outputs=navigation_outputs ) if __name__ == "__main__": app.launch(share=True)