Create app.py

app.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import glob
+import io
+import os
+import random
+import struct
+from contextlib import contextmanager
+from html import escape
+
+import msgpack
+import streamlit as st
+import torch
+import tqdm
+from huggingface_hub import HfFileSystem
+from transformers import AutoTokenizer
+
+st.set_page_config(layout="wide")
+
+MODEL_NAME = os.environ.get("MODEL_NAME", "MonetLLM/monet-vd-1.4B-100BT-hf")
+CONTEXT_WINDOW = int(os.environ.get("CONTEXT_WINDOW", "12"))
+CANDIDATE_THRESHOLD = int(os.environ.get("CANDIDATE_THRESHOLD", "50"))
+
+HORIZONTAL_STYLE = """<style class="hide-element">
+    /* Hides the style container and removes the extra spacing */
+    .element-container:has(.hide-element) {
+        display: none;
+    }
+    /*
+        The selector for >.element-container is necessary to avoid selecting the whole
+        body of the streamlit app, which is also a stVerticalBlock.
+    */
+    div[data-testid="stVerticalBlock"]:has(> .element-container .horizontal-marker) {
+        display: flex;
+        flex-direction: row !important;
+        flex-wrap: wrap;
+        gap: 0.5rem;
+        align-items: baseline;
+    }
+    /* Buttons and their parent container all have a width of 704px, which we need to override */
+    div[data-testid="stVerticalBlock"]:has(> .element-container .horizontal-marker) div {
+        width: max-content !important;
+    }
+    /* Just an example of how you would style buttons, if desired */
+    /*
+    div[data-testid="stVerticalBlock"]:has(> .element-container .horizontal-marker) button {
+        border-color: red;
+    }
+    */
+</style>"""
+
+
+@st.cache_resource
+def prepare_routing_resources():
+    fs = HfFileSystem()
+    for filename in fs.glob(f"datasets/{MODEL_NAME}-viewer-data/*"):
+        if not os.path.exists(os.path.basename(filename)):
+            print(f"[*] Download {filename}...")
+            fs.download(filename, ".")
+
+    input_tokens = torch.load("inputs.pt")
+
+    examples_tables = []
+    for i in tqdm.trange(len(glob.glob("examples-*.msgpack"))):
+        with open(f"examples-{i}.msgpack", "rb") as fp:
+            fp.seek(-4, io.SEEK_END)
+            table_size = struct.unpack(">I", fp.read(4))[0]
+
+            fp.seek(-(table_size + 4), io.SEEK_END)
+            examples_tables.append(msgpack.Unpacker(fp).unpack())
+
+    candidates = []
+    for i, table in enumerate(tqdm.tqdm(examples_tables)):
+        candidates.append([])
+        with open(f"examples-{i}.msgpack", "rb") as fp:
+            unpacker = msgpack.Unpacker(fp)
+            for j in range(len(table)):
+                if len(unpacker.unpack()) > CANDIDATE_THRESHOLD:
+                    candidates[-1].append(j)
+
+    routing_tables = []
+    for i in tqdm.trange(len(examples_tables)):
+        with open(f"routings-{i}.msgpack", "rb") as fp:
+            fp.seek(-4, io.SEEK_END)
+            table_size = struct.unpack(">I", fp.read(4))[0]
+
+            fp.seek(-(table_size + 4), io.SEEK_END)
+            routing_tables.append(msgpack.Unpacker(fp).unpack())
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    return input_tokens, examples_tables, routing_tables, candidates, tokenizer
+
+
+input_tokens, examples_tables, routing_tables, candidates, tokenizer = (
+    prepare_routing_resources()
+)
+
+
+def render_routing_examples_in_html(router_index: int, expert_id: int) -> str:
+    with open(f"examples-{router_index}.msgpack", "rb") as fp:
+        fp.seek(examples_tables[router_index][expert_id])
+        examples = msgpack.Unpacker(fp).unpack()
+    with open(f"routings-{router_index}.msgpack", "rb") as fp:
+        table = []
+        for i, j, _ in examples:
+            start = max(j - CONTEXT_WINDOW, 0)
+            end = min(j + CONTEXT_WINDOW, len(routing_tables[router_index][i]))
+
+            fp.seek(routing_tables[router_index][i][start])
+            unpacker = msgpack.Unpacker(fp, strict_map_key=False)
+            activated = [unpacker.unpack().get(expert_id, 0) for _ in range(start, end)]
+
+            full_text = tokenizer.decode(input_tokens[i])
+            encodings = tokenizer(full_text, add_special_tokens=False)
+            offset = len(encodings.input_ids) - input_tokens.size(1)
+
+            spans, lslice = [], None
+            for k in range(start, end):
+                if offset + k >= 0 and (sslice := encodings.token_to_chars(offset + k)):
+                    span, score = full_text[slice(*sslice)], activated[k - start]
+                    if lslice == sslice:
+                        score = max(spans.pop(-1)[1], score)
+                    spans.append((escape(span), score))
+                    lslice = sslice
+
+            spans = [
+                f"<span style='background-color: rgba(144, 238, 144, {score}' title='Routing: {score*100:.2f}%'>{span}</span>"
+                for span, score in spans
+            ]
+            table.append(
+                f"""
+                <tr>
+                    <td align='right'>
+                        <span style='font-weight: bold'>
+                            {escape(tokenizer.decode(input_tokens[i, j]))} ({activated[j - start] * 100:.2f}%)
+                        </span>
+                    </td>
+                    <td align='left'>
+                        (...) {"".join(spans)} (...)
+                    </td>
+                    <td align='right'>
+                        ({i}, {j})
+                    </td>
+                </tr>
+                """
+            )
+
+    return f"""
+        <div style='background-color: white; color: black; padding: 1em 3em; font-size: 12pt'>
+            <h2 style='font-size: 18pt'> Activated Examples of Group {router_index} / Expert {expert_id} </h2>
+            <table>
+                {"".join(table)}
+            </table>
+        </div>
+        """
+
+
+@contextmanager
+def st_horizontal():
+    st.markdown(HORIZONTAL_STYLE, unsafe_allow_html=True)
+    with st.container():
+        st.markdown(
+            '<span class="hide-element horizontal-marker"></span>',
+            unsafe_allow_html=True,
+        )
+        yield
+
+
+col1, col2 = st.columns(2)
+with col1:
+    router_groups = [f"Routing Group {i}" for i in range(len(examples_tables))]
+    router_index = st.selectbox("Expert Routing Group", router_groups, index=4)
+with col2:
+    expert_id = st.number_input("Expert Index", 0, len(examples_tables[0]), 54136)
+
+with st_horizontal():
+    show_btn = st.button("Show")
+    random_btn = st.button("Random")
+
+if show_btn or random_btn:
+    router_index = router_groups.index(router_index)
+    if random_btn:
+        expert_id = random.choice(candidates[router_index])
+    st.html(render_routing_examples_in_html(router_index, expert_id))