three additional 1B to 3B params models

app.py

@@ -1,5 +1,16 @@
 import os
 
+# Questions for Gradio
+# - Chat share button is enabled by default but thrown an error when clicked.
+# - How to add local images in HTML? (https://github.com/gradio-app/gradio/issues/884)
+# - How to allow Chatbot to fill the vertical space? (https://github.com/gradio-app/gradio/issues/4001)
+# TODO:
+# - Add the 1MB models, keras/gemma_1.1_instruct_7b_en
+# - Add retry button, for each model individually
+# - Add ability to route a message to a single model only.
+# - log_applied_layout_map: make it work for Llama3CausalLM and LlamaCausalLM (vicuna)
+# - display context length
+
 os.environ["KERAS_BACKEND"] = "jax"
 
 import gradio as gr
@@ -95,7 +106,7 @@ def bot_icon_select(model_name):
     if "gemma" in model_name:
         return "img/gemma.png"
     elif "llama" in model_name:
-        return "img/llama.png"
+        return "img/meta.png"
     elif "vicuna" in model_name:
         return "img/vicuna.png"
     elif "mistral" in model_name:
@@ -167,7 +178,7 @@ with gr.Blocks(fill_width=True, title="Keras demo") as demo:
             show_fullscreen_button=False,
             show_share_button=False,
             interactive=False,
-            scale=0.01,
+            scale=0,
             container=False,
         )
         gr.HTML(

models.py

@@ -2,11 +2,17 @@ import keras
 import keras_hub
 
 model_presets = [
+    # 8B params models
     "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",
+    # "keras/gemma_1.1_instruct_7b_en", # won't fit?
+    # 1-3B params models
+    "hf://meta-llama/Llama-3.2-1B-Instruct",
+    "hf://google/gemma-2b-it-keras",
+    "hf://meta-llama/Llama-3.2-3B-Instruct",
 ]
 
 model_labels = map(lambda s: s.removeprefix("hf://"), model_presets)
@@ -33,18 +39,27 @@ def get_default_layout_map(preset_name, device_mesh):
         or "mistral" in preset_name
         or "vicuna" in preset_name
     ):
-        return keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
+        layout_map = keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
+        # This line is missing for some Llama models (TODO: fix this in keras_hub)
+        layout_map["token_embedding/reverse_embeddings"] = ("batch", "model")
+        return layout_map
     elif "gemma" in preset_name:
         return keras_hub.models.GemmaBackbone.get_layout_map(device_mesh)
 
 
 def log_applied_layout_map(model):
+    print("Model class:", type(model).__name__)
+
     if "Gemma" in type(model).__name__:
         transformer_decoder_block_name = "decoder_block_1"
-    else:  # works for Llama, Mistral, Vicuna
+    elif "Llama" in type(model).__name__:  # works for Llama (Vicuna) and Llama3
+        transformer_decoder_block_name = "transformer_layer_1"
+    elif "Mistral" in type(model).__name__:
         transformer_decoder_block_name = "transformer_layer_1"
+    else:
+        print("Unknown architecture. Cannot display the applied layout.")
+        return
 
-    print("Model class:", type(model).__name__)
     # See how layer sharding was applied
     embedding_layer = model.backbone.get_layer("token_embedding")
     print(embedding_layer)
@@ -96,10 +111,3 @@ def load_model(preset):
 
     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