Hugging Face's logo

Spaces:
huggingface
/
keras-chatbot-arena
Running on TPU v5e

App Files Community
1
keras-chatbot-arena
File size: 4,981 Bytes 
2ca0c5e
 
 
 
40912b5
2ca0c5e
 
 
 
 
40912b5
 
 
 
 
2ca0c5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40912b5
1365804
40912b5
 
 
1365804
2ca0c5e
1365804
 
 
 
 
 
 
 
 
 
 
 
 
 
38f8411
 
2ca0c5e
 
 
40912b5
 
a2b7758
2ca0c5e
40912b5
 
 
2ca0c5e
40912b5
 
 
a2b7758
 
 
 
 
 
 
 
 
 
 
 
 
2ca0c5e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
 
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)
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
    ):
        layout_map = keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
        # Default layout map patch:
        # 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:
        layout_map = keras_hub.models.GemmaBackbone.get_layout_map(device_mesh)

        if "gemma-2b-" in preset_name:
            # Default layout map patch:
            # Gemma QKV weigts are shaped [NB_HEADS, EMBED_DIM, INNER_DIM]
            # Llama QKV weights are shaped [EMBED_DIM, NB_HEADS, INNER_DIM]
            # However:
            # The default layout map for KQV weights on Gemma is: (model_dim,data_dim,None)
            # Which means sharding NB_HEADS on the "model" dimension.
            # But gemma-2b-it-keras has only 1 head so this won't work: must patch it
            # TODO: fix this in the Gemma layout map in Keras hub.
            patch_key = "decoder_block.*attention.*(query|key|value).kernel"
            layout_map.pop(patch_key)
            layout_map[patch_key] = (None, "model", "batch")

        return layout_map


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"
    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

    # 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