app.py

import os
from typing import Dict, Tuple
from uuid import UUID

import altair as alt
import argilla as rg
from argilla.feedback import FeedbackDataset
from argilla.client.feedback.dataset.remote.dataset import RemoteFeedbackDataset
import gradio as gr
import pandas as pd


def obtain_source_target_datasets() -> (
    Tuple[
        FeedbackDataset | RemoteFeedbackDataset, FeedbackDataset | RemoteFeedbackDataset
    ]
):
    """
    This function returns the source and target datasets to be used in the application.

    Returns:
        A tuple with the source and target datasets. The source dataset is filtered by the response status 'pending'.

    """

    # Obtain the public dataset and see how many pending records are there
    source_dataset = rg.FeedbackDataset.from_argilla(
        os.getenv("SOURCE_DATASET"), workspace=os.getenv("SOURCE_WORKSPACE")
    )
    filtered_source_dataset = source_dataset.filter_by(response_status=["pending"])

    # Obtain a list of users from the private workspace
    target_dataset = rg.FeedbackDataset.from_argilla(
        os.getenv("RESULTS_DATASET"), workspace=os.getenv("RESULTS_WORKSPACE")
    )

    return filtered_source_dataset, target_dataset


def get_user_annotations_dictionary(
    dataset: FeedbackDataset | RemoteFeedbackDataset,
) -> Dict[str, int]:
    """
    This function returns a dictionary with the username as the key and the number of annotations as the value.

    Args:
        dataset: The dataset to be analyzed.
    Returns:
        A dictionary with the username as the key and the number of annotations as the value.
    """
    output = {}
    for record in dataset:
        for response in record.responses:
            if str(response.user_id) not in output.keys():
                output[str(response.user_id)] = 1
            else:
                output[str(response.user_id)] += 1

    # Changing the name of the keys, from the id to the username
    for key in list(output.keys()):
        output[rg.User.from_id(UUID(key)).username] = output.pop(key)

    return output


def donut_chart2() -> alt.Chart:
    # Load your data
    source_dataset, results = obtain_source_target_datasets()
    pending_records = len(source_dataset)
    annotated_records = len(results)

    # Prepare data for the donut chart
    source = pd.DataFrame({
        "values": [annotated_records, pending_records],
        "category": ["Completed", "Remaining"],
        "colors": ["#4CAF50", "#757575"]  # Green for Completed, Grey for Remaining
    })

    # Base chart for donut segments
    base = alt.Chart(source).encode(
        theta=alt.Theta(field="values", type="quantitative", stack=True),  # The angle encoding
        color=alt.Color(field="colors", type="nominal", legend=alt.Legend(title="Category"), scale=None),  # Direct color specification
        tooltip=['category', 'values']  # Tooltips for interactivity
    )

    # Arc marks for donut segments
    arcs = base.mark_arc(innerRadius=100, outerRadius=150, stroke="#fff")  # Increased inner radius for a thicker donut

    # Text marks for labels, adjusted to be placed within the donut segments
    text = base.mark_text(radiusOffset=-70, size=16, color='white').encode(
        text=alt.Text(field="values", type="quantitative"),  # The text to display
    )

    # Combine the arcs and text
    chart = arcs + text

    # Configure the view to have a transparent background
    chart = chart.configure_view(
        strokeWidth=0,  # Remove border around chart
        fill=None  # Ensure background is transparent
    ).configure_axis(
        grid=False  # Turn off the grid
    )

    return chart



def donut_chart() -> alt.Chart:
    # Load your data
    source_dataset, results = obtain_source_target_datasets()
    pending_records = len(source_dataset)
    annotated_records = len(results)

    # Prepare data for the donut chart
    source = pd.DataFrame({
        "values": [annotated_records, pending_records],
        "category": ["Completed", "Remaining"],
        "colors": ["#4CAF50", "#757575"]  # Green for Completed, Grey for Remaining
    })

    base = alt.Chart(source).encode(
        theta=alt.Theta("values:Q", stack=True),
        radius=alt.Radius(
            "values", scale=alt.Scale(type="sqrt", zero=True, rangeMin=20)
        ),
        color=alt.Color("category:N", legend=alt.Legend(title="Category")),
    )

    c1 = base.mark_arc(innerRadius=20, stroke="#fff")

    c2 = base.mark_text(radiusOffset=10).encode(text="values:Q")

    chart = c1 + c2

    return chart




def kpi_chart() -> alt.Chart:
    """
    This function returns a KPI chart with the total amount of annotators.

    Returns:
        An altair chart with the KPI chart.
    """

    # Obtain the total amount of annotators
    _, target_dataset = obtain_source_target_datasets()
    user_ids_annotations = get_user_annotations_dictionary(target_dataset)
    total_annotators = len(user_ids_annotations)

    # Assuming you have a DataFrame with user data, create a sample DataFrame
    data = pd.DataFrame({"Category": ["Total Contributors"], "Value": [total_annotators]})

    # Create Altair chart
    chart = (
        alt.Chart(data)
        .mark_text(fontSize=100, align="center", baseline="middle", color="steelblue")
        .encode(text="Value:N")
        .properties(title="Number of Contributors", width=250, height=200)
    )

    return chart


def obtain_top_5_users(user_ids_annotations: Dict[str, int]) -> pd.DataFrame:
    """
    This function returns the top 5 users with the most annotations.

    Args:
        user_ids_annotations: A dictionary with the user ids as the key and the number of annotations as the value.

    Returns:
        A pandas dataframe with the top 5 users with the most annotations.
    """

    dataframe = pd.DataFrame(
        user_ids_annotations.items(), columns=["Name", "Submitted Responses"]
    )
    dataframe = dataframe.sort_values(by="Submitted Responses", ascending=False)
    return dataframe.head(10)


def main() -> None:

    # Connect to the space with rg.init()
    rg.init(
        api_url=os.getenv("ARGILLA_API_URL"),
        api_key=os.getenv("ARGILLA_API_KEY"),
        extra_headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},
    )

    source_dataset, target_dataset = obtain_source_target_datasets()
    user_ids_annotations = get_user_annotations_dictionary(target_dataset)

    top5_dataframe = obtain_top_5_users(user_ids_annotations)

    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # 🗣️ The Prompt Collective Dashboad

            This Gradio dashboard shows the progress of the first "Data is Better Together" initiative to understand and collect good quality and diverse prompt for the OSS AI community.
            If you want to contribute to OSS AI, join [the Prompt Collective HF Space](https://huggingface.co/spaces/DIBT/prompt-collective).
            """
        )
        gr.Markdown(
            """
            ## 🚀 Contributors Progress

            How many records have been submitted, how many are still pending?
            """
        )
        plot = gr.Plot(label="Plot")
        demo.load(
            donut_chart2,
            inputs=[],
            outputs=[plot],
        )

        gr.Markdown(
            """
            ## 👾 Contributors Hall of Fame
            The number of all contributors and the top 10 contributors:
            """
        )

        with gr.Row():

            plot2 = gr.Plot(label="Plot")
            demo.load(
                kpi_chart,
                inputs=[],
                outputs=[plot2],
            )

            gr.Dataframe(
                value=top5_dataframe,
                headers=["Name", "Submitted Responses"],
                datatype=[
                    "str",
                    "number",
                ],
                row_count=10,
                col_count=(2, "fixed"),
                interactive=False,
            ),

    # Launch the Gradio interface
    demo.launch()


if __name__ == "__main__":
    main()