initial commit

.gitignore

@@ -0,0 +1,3 @@
+.DS_Store
+.vscode
+__pycache__

app.py

@@ -1,63 +1,220 @@
+import os
+
+os.environ["KERAS_BACKEND"] = "jax"
+
 import gradio as gr
-from huggingface_hub import InferenceClient
+from gradio import ChatMessage
+import keras_hub
+
+from chatstate import ChatState
+from models import (
+    model_presets,
+    load_model,
+    model_labels,
+    preset_to_website_url,
+    get_appropriate_chat_template,
+)
+
+model_labels_list = list(model_labels)
+
+# lod a warm up (compile) all the models
+models = []
+for preset in model_presets:
+    model = load_model(preset)
+    chat_template = get_appropriate_chat_template(preset)
+    chat_state = ChatState(model, "", chat_template)
+    prompt, response = chat_state.send_message("Hello")
+    print("model " + preset + "loaded and initialized.")
+    print("The model responded: " + response)
+
+models = [load_model(preset) for preset in model_presets]
+# model = keras_hub.models.Llama3CausalLM.from_preset(
+#     "hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16"
+# )
+# models = [model, model]
+
+
+def chat_turn_assistant_1(
+    model,
+    message,
+    history,
+    system_message,
+    preset,
+    # max_tokens,
+    # temperature,
+    # top_p,
+):
+    chat_template = get_appropriate_chat_template(preset)
+    chat_state = ChatState(model, system_message, chat_template)
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+    for msg in history:
+        msg = ChatMessage(**msg)
+        if msg.role == "user":
+            chat_state.add_to_history_as_user(msg.content)
+        elif msg.role == "assistant":
+            chat_state.add_to_history_as_model(msg.content)
 
+    prompt, response = chat_state.send_message(message)
+    history.append(ChatMessage(role="assistant", content=response))
+    return history
 
-def respond(
+
+def chat_turn_assistant(
     message,
-    history: list[tuple[str, str]],
+    sel1,
+    history1,
+    sel2,
+    history2,
     system_message,
-    max_tokens,
-    temperature,
-    top_p,
+    # max_tokens,
+    # temperature,
+    # top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
+    history1 = chat_turn_assistant_1(
+        models[sel1], message, history1, system_message, model_presets[sel1]
+    )
+    history2 = chat_turn_assistant_1(
+        models[sel2], message, history2, system_message, model_presets[sel2]
+    )
+    return "", history1, history2
+
+
+def chat_turn_user_1(message, history):
+    history.append(ChatMessage(role="user", content=message))
+    return history
+
+
+def chat_turn_user(message, history1, history2):
+    history1 = chat_turn_user_1(message, history1)
+    history2 = chat_turn_user_1(message, history2)
+    return "", history1, history2
+
+
+def bot_icon_select(model_name):
+    if "gemma" in model_name:
+        return "img/gemma.png"
+    elif "llama" in model_name:
+        return "img/llama.png"
+    elif "vicuna" in model_name:
+        return "img/vicuna.png"
+    elif "mistral" in model_name:
+        return "img/mistral.png"
+    # default
+    return "img/bot.png"
+
+
+def instantiate_chatbots(sel1, sel2):
+    model_name1 = model_presets[sel1]
+    chatbot1 = gr.Chatbot(
+        type="messages",
+        show_label=False,
+        avatar_images=("img/usr.png", bot_icon_select(model_name1)),
+    )
+    model_name2 = model_presets[sel2]
+    chatbot2 = gr.Chatbot(
+        type="messages",
+        show_label=False,
+        avatar_images=("img/usr.png", bot_icon_select(model_name2)),
+    )
+    return chatbot1, chatbot2
+
+
+def instantiate_select_boxes(sel1, sel2, model_labels):
+    sel1 = gr.Dropdown(
+        choices=[(name, i) for i, name in enumerate(model_labels)],
+        show_label=False,
+        info="<span style='color:black'>Selected model 1:</span> "
+        + "<a href='"
+        + preset_to_website_url(model_presets[sel1])
+        + "'>"
+        + preset_to_website_url(model_presets[sel1])
+        + "</a>",
+        value=sel1,
+    )
+    sel2 = gr.Dropdown(
+        choices=[(name, i) for i, name in enumerate(model_labels)],
+        show_label=False,
+        info="<span style='color:black'>Selected model 2:</span> "
+        + "<a href='"
+        + preset_to_website_url(model_presets[sel2])
+        + "'>"
+        + preset_to_website_url(model_presets[sel2])
+        + "</a>",
+        value=sel2,
+    )
+    return sel1, sel2
+
+
+def instantiate_chatbots_and_select_boxes(sel1, sel2, model_labels):
+    chatbot1, chatbot2 = instantiate_chatbots(sel1, sel2)
+    sel1, sel2 = instantiate_select_boxes(sel1, sel2, model_labels)
+    return sel1, chatbot1, sel2, chatbot2
+
+
+with gr.Blocks(fill_width=True, title="Keras demo") as demo:
+
+    with gr.Row():
+        gr.Image(
+            "img/keras_logo_k.png",
+            width=80,
+            height=80,
+            min_width=80,
+            show_label=False,
+            show_download_button=False,
+            show_fullscreen_button=False,
+            interactive=False,
+            scale=0.01,
+            container=False,
+        )
+        gr.HTML(
+            "<H2> Battle of the Keras chatbots on TPU</H2>"
+            + "All the models are loaded into the TPU memory. "
+            + "You can call them at will and compare their answers. <br/>"
+            + "The entire chat history is fed to the models at every submission."
+            + "This demno is runnig on a Google TPU v5e 2x4 (8 cores).",
+        )
+    with gr.Row():
+        sel1, sel2 = instantiate_select_boxes(0, 1, model_labels_list)
+
+    with gr.Row():
+        chatbot1, chatbot2 = instantiate_chatbots(sel1.value, sel2.value)
+
+    msg = gr.Textbox(
+        label="Your message:",
+    )
+    with gr.Row():
+        gr.ClearButton([msg, chatbot1, chatbot2])
+        with gr.Accordion("Additional settings", open=False):
+            system_message = gr.Textbox(
+                label="Sytem prompt",
+                value="You are a helpful assistant and your name is Eliza.",
+            )
+
+    sel1.select(
+        lambda sel1, sel2: instantiate_chatbots_and_select_boxes(
+            sel1, sel2, model_labels_list
         ),
-    ],
-)
+        inputs=[sel1, sel2],
+        outputs=[sel1, chatbot1, sel2, chatbot2],
+    )
+
+    sel2.select(
+        lambda sel1, sel2: instantiate_chatbots_and_select_boxes(
+            sel1, sel2, model_labels_list
+        ),
+        inputs=[sel1, sel2],
+        outputs=[sel1, chatbot1, sel2, chatbot2],
+    )
+
+    msg.submit(
+        chat_turn_user,
+        inputs=[msg, chatbot1, chatbot2],
+        outputs=[msg, chatbot1, chatbot2],
+    ).then(
+        chat_turn_assistant,
+        [msg, sel1, chatbot1, sel2, chatbot2, system_message],
+        outputs=[msg, chatbot1, chatbot2],
+    )
 
 
 if __name__ == "__main__":

chatstate.py

@@ -0,0 +1,94 @@
+# chat helper
+class ChatState:
+
+    def __init__(self, model, system="", chat_template="auto"):
+        chat_template = (
+            type(model).__name__ if chat_template == "auto" else chat_template
+        )
+
+        if chat_template == "Llama3CausalLM":
+            self.__START_TURN_SYSTEM__ = (
+                "<|start_header_id|>system<|end_header_id|>\n\n"
+            )
+            self.__START_TURN_USER__ = (
+                "<|start_header_id|>user<|end_header_id|>\n\n"
+            )
+            self.__START_TURN_MODEL__ = (
+                "<|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+            self.__END_TURN_SYSTEM__ = "<|eot_id|>"
+            self.__END_TURN_USER__ = "<|eot_id|>"
+            self.__END_TURN_MODEL__ = "<|eot_id|>"
+            print("Using chat template for: Llama")
+        elif chat_template == "GemmaCausalLM":
+            self.__START_TURN_SYSTEM__ = ""
+            self.__START_TURN_USER__ = "<start_of_turn>user\n"
+            self.__START_TURN_MODEL__ = "<start_of_turn>model\n"
+            self.__END_TURN_SYSTEM__ = "\n"
+            self.__END_TURN_USER__ = "<end_of_turn>\n"
+            self.__END_TURN_MODEL__ = "<end_of_turn>\n"
+            print("Using chat template for: Gemma")
+        elif chat_template == "MistralCausalLM":
+            self.__START_TURN_SYSTEM__ = ""
+            self.__START_TURN_USER__ = "[INST]"
+            self.__START_TURN_MODEL__ = ""
+            self.__END_TURN_SYSTEM__ = "<s>"
+            self.__END_TURN_USER__ = "[/INST]"
+            self.__END_TURN_MODEL__ = "</s>"
+            print("Using chat template for: Mistral")
+        elif chat_template == "Vicuna":
+            self.__START_TURN_SYSTEM__ = ""
+            self.__START_TURN_USER__ = "USER: "
+            self.__START_TURN_MODEL__ = "ASSISTANT: "
+            self.__END_TURN_SYSTEM__ = "\n\n"
+            self.__END_TURN_USER__ = "\n"
+            self.__END_TURN_MODEL__ = "</s>\n"
+            print("Using chat template for : Vicuna")
+        else:
+            assert (0, "Unknown turn tags for this model class")
+
+        self.model = model
+        self.system = system
+        self.history = []
+
+    def add_to_history_as_user(self, message):
+        self.history.append(
+            self.__START_TURN_USER__ + message + self.__END_TURN_USER__
+        )
+
+    def add_to_history_as_model(self, message):
+        self.history.append(
+            self.__START_TURN_MODEL__ + message + self.__END_TURN_MODEL__
+        )
+
+    def get_history(self):
+        return "".join([*self.history])
+
+    def get_full_prompt(self):
+        prompt = self.get_history() + self.__START_TURN_MODEL__
+        if len(self.system) > 0:
+            prompt = (
+                self.__START_TURN_SYSTEM__
+                + self.system
+                + self.__END_TURN_SYSTEM__
+                + prompt
+            )
+        return prompt
+
+    def send_message(self, message):
+        """
+        Handles sending a user message and getting a model response.
+
+        Args:
+            message: The user's message.
+
+        Returns:
+            The model's response.
+        """
+        self.add_to_history_as_user(message)
+        prompt = self.get_full_prompt()
+        response = self.model.generate(
+            prompt, max_length=1024, strip_prompt=True
+        )
+        self.add_to_history_as_model(response)
+        return (message, response)

models.py

@@ -0,0 +1,105 @@
+import keras
+import keras_hub
+
+model_presets = [
+    "hf://google/gemma-2-instruct-9b-keras",
+    "hf://meta-llama/Llama-3.1-8B-Instruct",
+    "hf://google/codegemma-7b-it-keras",
+    "hf://keras/mistral_instruct_7b_en",
+    "hf://keras/vicuna_1.5_7b_en",
+]
+
+model_labels = map(lambda s: s.removeprefix("hf://"), model_presets)
+model_labels = map(lambda s: s.removeprefix("google/"), model_labels)
+model_labels = map(lambda s: s.removeprefix("keras/"), model_labels)
+model_labels = map(lambda s: s.removeprefix("meta-llama/"), model_labels)
+
+
+def preset_to_website_url(preset):
+    preset = preset.removeprefix("hf://")
+    url = "http://huggingface.co/" + preset
+    return url
+
+
+def get_appropriate_chat_template(preset):
+    return "Vicuna" if "vicuna" in preset else "auto"
+
+
+def get_default_layout_map(preset_name, device_mesh):
+    # Llama's default layout map works for mistral and vicuna
+    # because their transformer layers have the same names.
+    if (
+        "Llama" in preset_name
+        or "mistral" in preset_name
+        or "vicuna" in preset_name
+    ):
+        return keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
+    elif "gemma" in preset_name:
+        return keras_hub.models.GemmaBackbone.get_layout_map(device_mesh)
+
+
+def log_applied_layout_map(model):
+    if "Gemma" in type(model):
+        transformer_decoder_block_name = "decoder_block_1"
+    elif "Llama3" in type(model) or "Mistral" in type(model):
+        transformer_decoder_block_name = "transformer_layer_1"
+    else:
+        assert (0, "Model type not recognized. Cannot display model layout.")
+        # See how layer sharding was applied
+        embedding_layer = model.backbone.get_layer("token_embedding")
+        print(embedding_layer)
+        decoder_block = model.backbone.get_layer(transformer_decoder_block_name)
+        print(type(decoder_block))
+        for variable in embedding_layer.weights + decoder_block.weights:
+            print(
+                f"{variable.path:<58}  \
+                  {str(variable.shape):<16}  \
+                  {str(variable.value.sharding.spec):<35} \
+                  {str(variable.dtype)}"
+            )
+
+
+def load_model(preset):
+    devices = keras.distribution.list_devices()
+    device_mesh = keras.distribution.DeviceMesh(
+        shape=(1, len(devices)), axis_names=["batch", "model"], devices=devices
+    )
+    model_parallel = keras.distribution.ModelParallel(
+        layout_map=get_default_layout_map(preset, device_mesh),
+        batch_dim_name="batch",
+    )
+
+    with model_parallel.scope():
+        # These two buggy models need this workaround to be loaded in bfloat16
+        if "google/gemma-2-instruct-9b-keras" in preset:
+            model = keras_hub.models.GemmaCausalLM(
+                backbone=keras_hub.models.GemmaBackbone.from_preset(
+                    preset, dtype="bfloat16"
+                ),
+                preprocessor=keras_hub.models.GemmaCausalLMPreprocessor.from_preset(
+                    preset
+                ),
+            )
+        elif "meta-llama/Llama-3.1-8B-Instruct" in preset:
+            model = keras_hub.models.Llama3CausalLM(
+                backbone=keras_hub.models.Llama3Backbone.from_preset(
+                    preset, dtype="bfloat16"
+                ),
+                preprocessor=keras_hub.models.Llama3CausalLMPreprocessor.from_preset(
+                    preset
+                ),
+            )
+        else:
+            model = keras_hub.models.CausalLM.from_preset(
+                preset, dtype="bfloat16"
+            )
+
+    log_applied_layout_map(model)
+    return model
+
+
+# Some small models too
+# model1 = keras_hub.models.CausalLM.from_preset("hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16")
+# model2 = keras_hub.models.CausalLM.from_preset("hf://google/gemma-2b-it-keras", dtype="bfloat16")
+# model3 = keras_hub.models.CausalLM.from_preset("hf://meta-llama/Llama-3.2-3B-Instruct", dtype="bfloat16")
+# keras/gemma_1.1_instruct_7b_en

requirements.txt

@@ -1 +1,6 @@
-huggingface_hub==0.25.2
+--find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
+jax[tpu]
+keras>=3
+keras-hub
+safetensors
+huggingface_hub