app.py

import json
import logging
import os
import random
from functools import partial

import gradio as gr
from datasets import Dataset, load_dataset
from dotenv import load_dotenv
import tempfile
import base64
import vertexai
from vertexai.generative_models import GenerativeModel, Part, GenerationResponse
import vertexai.preview.generative_models as generative_models


def multiturn_generate_content(things) -> GenerationResponse:
    vertexai.init(project="134994101092", location="europe-west6")
    model = GenerativeModel(
        "projects/134994101092/locations/europe-west6/endpoints/1920706076236316672",
        system_instruction=[textsi_1],
    )
    chat = model.start_chat()
    json_objects = json.dumps({"things": things})
    response = chat.send_message(
        [
            json_objects
        ],
        generation_config=generation_config,
        safety_settings=safety_settings,
    )
    return response


textsi_1 = """You are a creative chatbot called MakerBot, built for Maker Faire Aarhus. You have been built by Aldan Creo, an AI researcher. Your job is to invent new things that can be built using three things. The user will give you a list of three things, in JSON format, and you need to write what you would build, and how you would build it. The names of the things can be multiple words. Commas indicate synonyms or different ways to call that thing. You should try to use the three things. If it is impossible to use the three things to build something, you must explicitly say that you have not been able to think about how to use that thing, and say what it is. You must answer in Danish. Your answer must follow the structure {\"What\": \"Navn på opfindelsen\", \"How\": \"Hvordan man bygger den ting ved hjælp af de tre objekter, som brugeren har givet.\"}."""

generation_config = {
    "max_output_tokens": 2048,
    "temperature": 1,
    "top_p": 1,
}

safety_settings = {
    generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
    generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_models.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
}


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

load_dotenv()

# Get the credentials to use the Google API from the env variable GOOGLE_APP_CREDENTIALS_JSON and save it to a temp file
with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as f:
    f.write(os.environ["GOOGLE_APP_CREDENTIALS_JSON"])
    f.flush()
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = f.name

# dataset = load_dataset("detection-datasets/coco")
it_dataset = (
    load_dataset("imagenet-1k", split="train", streaming=True, trust_remote_code=True)
    .shuffle(42)
    .skip(0)
    .take(1000)
)


def gen_from_iterable_dataset(iterable_ds):
    """
    Convert an iterable dataset to a generator
    """
    yield from iterable_ds


dataset = Dataset.from_generator(
    partial(gen_from_iterable_dataset, it_dataset), features=it_dataset.features
)

# imagenet_categories_data.json is a JSON file containing a hierarchy of ImageNet categories.
# We want to take all categories under "artifact, artefact".
# Each node has this structure:
# {
#     "id": 1,
#     "name": "entity",
#     "children": ...
# }
with open("imagenet_categories_data.json") as f:
    data = json.load(f)

    # Recursively find all categories under "artifact, artefact".
    # We want to get all the "index" values of the leaf nodes. Nodes that are not leaf nodes have a "children" key.
    def find_categories(node):
        if "children" in node:
            for child in node["children"]:
                yield from find_categories(child)
        elif "index" in node:
            yield node["index"]

    broad_categories = data["children"]
    artifact_category = next(
        filter(lambda x: x["name"] == "artifact, artefact", broad_categories)
    )
    artifact_categories = list(find_categories(artifact_category))
    # logger.info(f"Artifact categories: {artifact_categories}")


def filter_imgs_by_label(x):
    """
    Filter out the images that have label -1
    """
    return x["label"] in artifact_categories


dataset = dataset.filter(filter_imgs_by_label)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

load_dotenv()


def get_user_prompt():
    # Pick the first 3 images and labels
    images = []
    machine_labels = []
    human_labels = []
    for i in range(3):
        data = dataset[random.randint(0, len(dataset) - 1)]
        images.append(data["image"])
        # Get the label as a human readable string
        machine_labels.append(data["label"])
        human_label = dataset.features["label"].int2str(data["label"])
        human_labels.append(human_label)
    return {
        "images": images,
        "machine_labels": machine_labels,
        "human_labels": human_labels,
    }


hf_writer = gr.HuggingFaceDatasetSaver(
    hf_token=os.environ["HF_TOKEN"],
    dataset_name="acmc/maker-faire-bot",
    private=True,
)
csv_writer = gr.CSVLogger()

theme = gr.themes.Default(primary_hue="cyan", secondary_hue="fuchsia")

translation_table = {
    "Maker Faire Bot": "Maker Faire Bot",
    "Think about these objects...": "## Tænk på disse objekter...",
    "We want to build a Maker Faire Bot that can generate creative ideas. Help us by providing ideas on what you'd build with the following three objects!": "Vi vil bygge en Maker Faire Bot, der kan generere kreative ideer. Hjælp os ved at give ideer til, hvad du ville bygge med de følgende tre objekter!",
    "Change": "Skift",
    "What would you build with these 3 things?": "Hvad ville du bygge med disse 3 ting?",
    "For example, if you have a roll of string, a camera, and a loudspeaker, you could build an electronic guitar. If you can write in Danish, that's great!": "For eksempel, hvis du har en rulle snor, et kamera og en højttaler, kunne du bygge en elektronisk guitar. Hvis du kan skrive på dansk, er det fantastisk!",
    "It doesn't need to be a very long explanation, just a few sentences to help the bot understand your idea.": "Det behøver ikke være en meget lang forklaring, bare et par sætninger for at hjælpe robotten med at forstå din idé.",
    "Submit Your Answer": "Indsend dit svar",
    "New Prompt": "Ny opgave",
    "How would you build it?": "Hvordan ville du bygge det?",
    "This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our policy.": "Dette er et eksperimentelt projekt. Dine data er anonyme og vil blive brugt til at træne en AI-model. Ved at bruge dette værktøj accepterer du vores politik.",
    "(example): An digital electronic guitar": "(eksempel): En digital elektronisk guitar",
    """I would use the roll of string to create the strings of the guitar, and the camera to analyze the hand movements. Then, I would use an AI model to predict the chords and play the sound through the loudspeaker.""": """Jeg ville bruge snoren til at skabe guitarens strenge, og kameraet til at analysere håndbevægelserne. Derefter ville jeg bruge en AI-model til at forudsige akkorderne og afspille lyden gennem højttaleren.""",
    "Ask the Bot": "Spørg robotten",
    "The bot needs human ideas to learn - can you provide a new one?": "Robotten har brug for menneskelige ideer for at lære - kan du give en ny?",
    "This demo has been built by [Aldan Creo](https://acmc-website.web.app/) for [Maker Faire Aarhus](https://aarhus.makerfaire.com/). The images are from the [ImageNet](https://www.image-net.org/) dataset.": "Denne demo er lavet af [Aldan Creo](https://acmc-website.web.app/) til [Maker Faire Aarhus](https://aarhus.makerfaire.com/). Billederne er fra [ImageNet](https://www.image-net.org/)-datasættet.",
    "The explanation should have at least 20 characters.": "Forklaringen skal have mindst 20 tegn.",
    "Can you think of something better? **Teach the bot** by writing your own idea above!": "Kan du tænke på noget bedre? **Lær robotten** ved at skrive din egen idé ovenfor!",
    "This is an experimental project where we teach a bot to think like a Maker. You can choose three objects by clicking the 'Change' button under the images below. Then, you can **ask the bot what it would build with those objects** by clicking on the 'Ask the Bot' button. What will it come up with?": "Dette er et eksperimentelt projekt, hvor vi lærer en robot at tænke som en Maker. Du kan vælge tre objekter ved at klikke på 'Skift'-knappen under billederne nedenfor. Derefter kan du **spørge robotten, hvad den ville bygge med de objekter** ved at klikke på 'Spørg robotten'-knappen. Hvad vil den finde på?",
}


def get_bilingual_string(key):
    return f"{translation_table[key]}  //  {key}"


with gr.Blocks(theme=theme) as demo:
    with gr.Row() as header:
        gr.Image(
            "maker-faire-logo.webp",
            show_download_button=False,
            show_label=False,
            show_share_button=False,
            container=False,
            # height=100,
            scale=0.2,
        )
        gr.Markdown(
            get_bilingual_string("Maker Faire Bot"),
            visible=False,
        )

    user_prompt = gr.State(get_user_prompt())
    last_bot_response = gr.State({"What": "", "How": ""})

    gr.Markdown(get_bilingual_string("Think about these objects...") + " 🤔")
    gr.Markdown(
        get_bilingual_string(
            "We want to build a Maker Faire Bot that can generate creative ideas. Help us by providing ideas on what you'd build with the following three objects!"
        ) + " 🙌",
        visible=False,
    )
    gr.Markdown(
        get_bilingual_string(
            "This is an experimental project where we teach a bot to think like a Maker. You can choose three objects by clicking the 'Change' button under the images below. Then, you can **ask the bot what it would build with those objects** by clicking on the 'Ask the Bot' button. What will it come up with?"
        )
    )

    image_components = []
    with gr.Row(variant="panel") as row:

        def change_image(this_i, user_prompt):
            logger.info(
                f"Current user prompt: {user_prompt}, current image index: {this_i}"
            )
            data = dataset[random.randint(0, len(dataset) - 1)]
            new_user_prompt = user_prompt.copy()
            new_user_prompt["images"][this_i] = data["image"]
            new_user_prompt["machine_labels"][this_i] = data["label"]
            new_user_prompt["human_labels"][this_i] = dataset.features["label"].int2str(
                data["label"]
            )
            logger.info(f"New user prompt: {new_user_prompt}")
            return (
                new_user_prompt,
                new_user_prompt["images"][this_i],
                gr.update(
                    label=new_user_prompt["human_labels"][this_i],
                ),
            )

        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][0],
                label=user_prompt.value["human_labels"][0],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change") + " 🔁", variant="secondary")

            btn.click(
                lambda *args: change_image(0, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )
        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][1],
                label=user_prompt.value["human_labels"][1],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change") + " 🔁", variant="secondary")

            btn.click(
                lambda *args: change_image(1, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )
        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][2],
                label=user_prompt.value["human_labels"][2],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change") + " 🔁", variant="secondary")

            btn.click(
                lambda *args: change_image(2, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )

    def ask_the_bot(user_prompt):
        response = multiturn_generate_content(things=user_prompt["human_labels"])
        json_bot_response = json.loads(response.candidates[0].text)
        logger.info(f"Bot response: {json_bot_response}")
        return json_bot_response, json_bot_response["What"], json_bot_response["How"]

    ask_the_bot_btn = gr.Button(get_bilingual_string("Ask the Bot") + " 🤖", variant="primary")

    user_answer_object = gr.Textbox(
        placeholder=get_bilingual_string("(example): An digital electronic guitar"),
        label=get_bilingual_string("What would you build with these 3 things?") + " 🤔",
        info=get_bilingual_string(
            "For example, if you have a roll of string, a camera, and a loudspeaker, you could build an electronic guitar. If you can write in Danish, that's great!"
        ),
    )
    user_answer_explanation = gr.TextArea(
        label=get_bilingual_string("How would you build it?") + " 🛠️",
        # The example uses a roll of string, a camera, and a loudspeaker to build an electronic guitar.
        placeholder=get_bilingual_string(
            """I would use the roll of string to create the strings of the guitar, and the camera to analyze the hand movements. Then, I would use an AI model to predict the chords and play the sound through the loudspeaker."""
        ),
        info=get_bilingual_string(
            "It doesn't need to be a very long explanation, just a few sentences to help the bot understand your idea."
        ),
    )

    csv_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_csv",
    )
    hf_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_hf",
    )

    gr.Markdown(get_bilingual_string("Can you think of something better? **Teach the bot** by writing your own idea above!") + " 🧠")

    submit_btn = gr.Button(get_bilingual_string("Submit Your Answer") + " 📩", variant="primary")

    def log_results(prompt, object, explanation):
        # Is the last bot response the same as the object or explanation?
        # If so, we don't want to log it
        if (
            last_bot_response.value["What"] == object
            or last_bot_response.value["How"] == explanation
        ):
            raise gr.Error(get_bilingual_string("The bot needs human ideas to learn - can you provide a new one?"))

        # The description should have at least 20 characters
        if len(explanation) < 20:
            raise gr.Error(get_bilingual_string("The explanation should have at least 20 characters."))
        logger.info(f"logging - Prompt: {prompt}")
        # csv_writer.flag(
        #    [
        #        {
        #            "machine_labels": prompt["machine_labels"],
        #            "human_labels": prompt["human_labels"],
        #        },
        #        object,
        #        explanation,
        #    ]
        # )
        hf_writer.flag(
            [
                {
                    "machine_labels": prompt["machine_labels"],
                    "human_labels": prompt["human_labels"],
                },
                object,
                explanation,
            ]
        )
        return ["", ""]  # Clear the textboxes

    submit_btn.click(
        log_results,
        inputs=[user_prompt, user_answer_object, user_answer_explanation],
        outputs=[user_answer_object, user_answer_explanation],
        preprocess=True,
    )

    ask_the_bot_btn.click(ask_the_bot, inputs=[user_prompt], outputs=[last_bot_response, user_answer_object, user_answer_explanation], preprocess=True)

    # def renew_prompt(image_components):
    #    new_prompt = get_user_prompt()
    #    for i in range(len(new_prompt["images"])):
    #        image_components[i].update(
    #            url=new_prompt["images"][i],
    #            label=new_prompt["human_labels"][i],
    #        )
    #    return new_prompt

    # new_prompt_btn = gr.Button(get_bilingual_string("New Prompt"), variant="secondary")
    # new_prompt_btn.click(
    #    renew_prompt,
    #    inputs=image_components,
    #    outputs=[user_prompt],
    #    # preprocess=True,
    # )

    gr.Markdown(
        get_bilingual_string(
            "This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our policy."
        )
    )

    # Attribution information
    gr.Markdown(
        get_bilingual_string(
            "This demo has been built by [Aldan Creo](https://acmc-website.web.app/) for [Maker Faire Aarhus](https://aarhus.makerfaire.com/). The images are from the [ImageNet](https://www.image-net.org/) dataset."
        )
    )
if __name__ == "__main__":
    demo.launch()