app.py

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)