:tada: init

README.md

@@ -1,10 +1,11 @@
 ---
-title: Tokenvisor
-emoji: 👀
-colorFrom: green
+title: Tokenvisor SD
+emoji: 🥽
+
+colorFrom: red
 colorTo: green
 sdk: gradio
-sdk_version: 4.36.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -1,94 +1,141 @@
-import os
-
 import gradio as gr
-import pandas as pd
-from gradio.themes import colors
-from transformers import AutoTokenizer
+import torch
+from diffusers import DiffusionPipeline
+from transformers import AutoTokenizer, CLIPTokenizerFast, T5TokenizerFast
+
+
+def load_tokenizers(model_id: str) -> list[CLIPTokenizerFast | T5TokenizerFast | None]:
+    config = DiffusionPipeline.load_config(model_id)
+    num_tokenizers = sum('tokenizer' in key for key in config.keys())
 
-os.environ['TOKENIZERS_PARALLELISM'] = "false"
+    if not 1 <= num_tokenizers <= 3:
+        raise gr.Error(f'Invalid number of tokenizers: {num_tokenizers}')
 
+    tokenizers = [
+        AutoTokenizer.from_pretrained(model_id, subfolder=f'tokenizer{"" if i == 0 else f"_{i + 1}"}')
+        for i in range(num_tokenizers)
+    ]
 
-# Function to map tokenized text to IDs
-def inference(
-        text="",
-        model_id="openai/clip-vit-large-patch14",
-        progress=gr.Progress()
-) -> (list[str, str], list[str, str], pd.DataFrame):
-    if text == "":
-        return [], [], pd.DataFrame()
+    # Pad the list with None if there are fewer than 3 tokenizers
+    tokenizers.extend([None] * (3 - num_tokenizers))
 
-    progress(0, desc='Loading tokenizer...')
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    return tokenizers
 
-    # Use tokenizer to tokenize the text
-    progress(0.5, desc='Tokenizing text...')
-    text_inputs = tokenizer(text, return_tensors='pt')
 
-    input_ids = text_inputs['input_ids'].tolist()[0]  # Convert tensor to list
+@torch.inference_mode()
+def inference(model_id: str, input_text: str):
+    tokenizers = load_tokenizers(model_id)
 
-    # Create pairs of tokens and IDs
-    tokens = [tokenizer.decode([id_]) for id_ in input_ids]
-    token_pairs = []
+    text_pairs_components = []
+    special_tokens_components = []
+    for i, tokenizer in enumerate(tokenizers):
+        if tokenizer:
+            label_text = f'Tokenizer {i + 1}: {tokenizer.__class__.__name__}'
 
-    for token, id_ in zip(tokens, input_ids):
-        token_pairs.append((token, str(id_)))
+            # テキストとトークンIDのペアを作成
+            input_ids = tokenizer(
+                text=input_text,
+                truncation=True,
+                return_length=False,
+                return_overflowing_tokens=False
+            ).input_ids
+            decoded_tokens = [tokenizer.decode(id_) for id_ in input_ids]
+            token_pairs = [(str(token), str(id_)) for token, id_ in zip(decoded_tokens, input_ids)]
+            output_text_pair_component = gr.HighlightedText(
+                label=label_text,
+                value=token_pairs,
+                visible=True,
+                show_legend=True,
+            )
 
-    # Count the number of characters and tokens
-    pos_count = pd.DataFrame({
-        "Char Count": [len(text)],
-        "Token Count": [len(token_pairs)]
-    })
+            # スペシャルトークンを追加
+            special_tokens = []
+            for k, v in tokenizer.special_tokens_map.items():
+                if k == 'additional_special_tokens':
+                    continue
+                special_token_map = (str(k), str(v))
+                special_tokens.append(special_token_map)
+            output_special_tokens_component = gr.HighlightedText(
+                label=label_text,
+                value=special_tokens,
+                visible=True,
+                show_legend=True,
+            )
+        else:
+            output_text_pair_component = gr.HighlightedText(visible=False)
+            output_special_tokens_component = gr.HighlightedText(visible=False)
 
-    # Create list of special tokens
-    special_tokens = []
-    for k, v in tokenizer.special_tokens_map.items():
-        if k == 'additional_special_tokens':
-            continue
-        sp_token_map = [str(k), str(v)]
-        special_tokens.append(sp_token_map)
+        text_pairs_components.append(output_text_pair_component)
+        special_tokens_components.append(output_special_tokens_component)
 
-    return token_pairs, special_tokens, pos_count
+    return text_pairs_components + special_tokens_components
 
 
 if __name__ == '__main__':
-    iface = gr.Interface(
-        fn=inference,
-        inputs=[
-            gr.Textbox(label="Text"),
-            gr.Dropdown(
-                label="Model",
+    theme = gr.themes.Soft(
+        primary_hue=gr.themes.colors.emerald,
+        secondary_hue=gr.themes.colors.emerald,
+    )
+    with gr.Blocks(theme=theme) as demo:
+        with gr.Column():
+            input_model_id = gr.Dropdown(
+                label='Model ID',
                 choices=[
-                    "openai/clip-vit-large-patch14",
-                    "google/gemma-7b",
-                    "google-bert/bert-base-uncased",
-                    "google/flan-t5-base",
-                    "openai-community/gpt2",
-                    "rinna/japanese-gpt-1b",
-                    "cyberagent/open-calm-7b",
+                    'black-forest-labs/FLUX.1-dev',
+                    'black-forest-labs/FLUX.1-schnell',
+                    'stabilityai/stable-diffusion-3-medium-diffusers',
+                    'stabilityai/stable-diffusion-xl-base-1.0',
+                    'stable-diffusion-v1-5/stable-diffusion-v1-5',
+                    'stabilityai/japanese-stable-diffusion-xl',
+                    'rinna/japanese-stable-diffusion',
                 ],
-                value="openai/clip-vit-large-patch14"
-            ),
-        ],
-        outputs=[
-            gr.Highlightedtext(label="Highlighted Text"),
-            gr.Highlightedtext(label="Special Tokens", combine_adjacent=True, adjacent_separator=' / '),
-            gr.Dataframe(label="Position Count"),
-        ],
-        examples=[
-            ["When I told my computer I needed a break, it froze.", "openai/clip-vit-large-patch14"],
-            ["Yesterday, I thought my cat was studying for her degree in philosophy because she sat on my book, "
-             "but turns out she was just trying to hatch a plot to steal my dinner.", "openai/clip-vit-large-patch14"],
-            ["The square root of x is the cube root of y. What is y to the power of 2, if x = 4?",
-             "google/flan-t5-base"],
-            ["In my home country, it's a custom to say 'いただきマサチューセッツ' before we start eating a meal.",
-             "google/gemma-7b"],
-            ["日本で一番高い山は富士山ですが、二番目に高い山は何ですか？", "rinna/japanese-gpt-1b"],
-        ],
-        cache_examples=True,
-        title="TokenVisor 👀",
-        description="Visualize how the Tokenizer used in Hugging Face's Transformers library tokenizes text.",
-        theme=gr.Theme(primary_hue=colors.green, secondary_hue=colors.yellow),
-        allow_flagging="never",
+                value='black-forest-labs/FLUX.1-dev',
+            )
+            input_text = gr.Textbox(
+                label='Input Text',
+                placeholder='Enter text here',
+            )
 
-    )
-    iface.queue().launch()
+            with gr.Tab(label='Tokenization Outputs'):
+                with gr.Column():
+                    output_highlighted_text_1 = gr.HighlightedText()
+                    output_highlighted_text_2 = gr.HighlightedText()
+                    output_highlighted_text_3 = gr.HighlightedText()
+            with gr.Tab(label='Special Tokens'):
+                with gr.Column():
+                    output_special_tokens_1 = gr.HighlightedText()
+                    output_special_tokens_2 = gr.HighlightedText()
+                    output_special_tokens_3 = gr.HighlightedText()
+
+            with gr.Row():
+                clear_button = gr.ClearButton(components=[input_text])
+                submit_button = gr.Button('Run', variant='primary')
+
+        all_inputs = [input_model_id, input_text]
+        all_output = [
+            output_highlighted_text_1,
+            output_highlighted_text_2,
+            output_highlighted_text_3,
+            output_special_tokens_1,
+            output_special_tokens_2,
+            output_special_tokens_3,
+        ]
+        submit_button.click(
+            fn=inference,
+            inputs=all_inputs,
+            outputs=all_output
+        )
+
+        examples = gr.Examples(
+            fn=inference,
+            inputs=all_inputs,
+            outputs=all_output,
+            examples=[
+                ['black-forest-labs/FLUX.1-dev', 'a photo of cat'],
+                ['stabilityai/stable-diffusion-3-medium-diffusers', 'cat holding sign saying "I am a cat"'],
+                ['rinna/japanese-stable-diffusion', '空を飛んでいるネコの写真　油絵']
+            ],
+            cache_examples=True
+        )
+
+    demo.queue().launch()

requirements.txt

@@ -3,4 +3,5 @@ transformers
 safetensors
 accelerate
 diffusers
-sentencepiece
+sentencepiece
+protobuf