import json
import gradio as gr
from datasets import load_dataset

ds = load_dataset("tti-bias/identities", split="train")

clusters_12 = json.load(open("clusters/id_all_blip_clusters_12.json"))
clusters_24 = json.load(open("clusters/id_all_blip_clusters_12.json"))
clusters_48 = json.load(open("clusters/id_all_blip_clusters_48.json"))

clusters_by_size = {
    12: clusters_12,
    24: clusters_24,
    48: clusters_48,
}

def show_cluster(cl_id, num_clusters):
    cl_dct = clusters_by_size[num_clusters][cl_id]
    return len(cl_dct['img_path_list']), cl_dct["labels_gender"], cl_dct["labels_model"], cl_dct["labels_ethnicity"], cl_dct["img_path_list"][:5]

demo = gr.Interface(fn=show_cluster, inputs=["text", gr.Radio([12, 24, 48])], outputs=["text", "text", "text", "text", "text"])