Update app.py

app.py

@@ -22,26 +22,54 @@ def get_available_models() -> list[str]:
 
 
 def tokenize_text(
-    builtin_model: str, custom_model: str | None, text: str
-) -> tuple[str | None, int, int]:
+    model_name: str, text: str
+) -> tuple[str | None, str | None, int | None, dict | None, int, int]:
     """处理tokenize请求"""
-    if not builtin_model:
-        return "Please choose a model and input some texts", 0, 0
+    if not model_name:
+        return "Please choose a model and input some texts", None, None, None, 0, 0
     if not text:
         text = "Please choose a model and input some texts"
 
     try:
         # 加载tokenizer
-        if custom_model:
+        model_path = os.path.join("models", model_name)
+        if os.path.isdir(model_path):
             tokenizer = AutoTokenizer.from_pretrained(
-                custom_model, trust_remote_code=True, device_map="cpu"
+                model_path, trust_remote_code=True, device_map="cpu"
             )
         else:
-            model_path = os.path.join("models", builtin_model)
             tokenizer = AutoTokenizer.from_pretrained(
-                model_path, trust_remote_code=True, device_map="cpu"
+                model_name, trust_remote_code=True, device_map="cpu"
             )
 
+        tokenizer_type = tokenizer.__class__.__name__
+
+        if hasattr(tokenizer, "vocab_size"):
+            vocab_size = tokenizer.vocab_size
+        elif hasattr(tokenizer, "get_vocab"):
+            vocab_size = len(tokenizer.get_vocab())
+        else:
+            vocab_size = -1
+
+        sp_token_list = [
+            "pad_token",
+            "eos_token",
+            "bos_token",
+            "sep_token",
+            "cls_token",
+            "unk_token",
+            "mask_token",
+        ]
+        special_tokens = {}
+        for token_name in sp_token_list:
+            if (
+                hasattr(tokenizer, token_name)
+                and getattr(tokenizer, token_name) is not None
+            ):
+                token_value = getattr(tokenizer, token_name)
+                if token_value and str(token_value).strip():
+                    special_tokens[token_name] = str(token_value)
+
         # Tokenize处理
         input_ids = tokenizer.encode(text, add_special_tokens=True)
 
@@ -68,11 +96,18 @@ def tokenize_text(
         token_len = len(input_ids)
         char_len = len(text)
 
-        return "".join(html_parts), token_len, char_len
+        return (
+            "".join(html_parts),
+            tokenizer_type,
+            vocab_size,
+            special_tokens,
+            token_len,
+            char_len,
+        )
 
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        return error_msg, 0, 0
+        return error_msg, None, None, None, 0, 0
 
 
 banner_md = """# 🎨 Tokenize it!
@@ -82,16 +117,18 @@ Powerful token visualization tool for your text inputs. 🚀
 Works for LLMs both online and *locally* on your machine!"""
 banner = gr.Markdown(banner_md)
 model_selector = gr.Dropdown(
-    label="Built-in Model", choices=get_available_models(), interactive=True
-)
-custom_model = gr.Textbox(
-    label="Custom Model",
-    placeholder="Enter your custom model name. e.g. Qwen/QwQ-32B. To use built-in models, please keep this EMPTY!",
-    lines=1,
+    label="Choose or enter model name",
+    choices=get_available_models(),
+    interactive=True,
+    allow_custom_value=True,
 )
 text_input = gr.Textbox(label="Input Text", placeholder="Hello World!", lines=4)
 submit_btn = gr.Button("🚀 Tokenize!", variant="primary")
 
+tokenizer_type = gr.Textbox(label="Tokenizer Type", interactive=False)
+vocab_size = gr.Number(label="Vocab Size", interactive=False)
+sp_tokens = gr.JSON(label="Special Tokens")
+
 output_html = gr.HTML(label="Tokenized Output", elem_classes="token-output")
 token_count = gr.Number(label="Token Count", value=0, interactive=False)
 char_count = gr.Number(label="Character Count", value=0, interactive=False)
@@ -99,19 +136,24 @@ char_count = gr.Number(label="Character Count", value=0, interactive=False)
 with gr.Blocks(title="Token Visualizer", theme="NoCrypt/miku") as webui:
     banner.render()
 
-    with gr.Column():
-        with gr.TabItem("Built-in Model"):
-            model_selector.render()
-        with gr.TabItem("Custom Model"):
-            custom_model.render()
-        text_input.render()
-        submit_btn.render()
-
     with gr.Column():
         with gr.Row():
-            token_count.render()
-            char_count.render()
-        output_html.render()
+            with gr.Column():
+                model_selector.render()
+                text_input.render()
+                submit_btn.render()
+            with gr.Column():
+                with gr.Accordion("Details", open=False):
+                    with gr.Row():
+                        tokenizer_type.render()
+                        vocab_size.render()
+                    sp_tokens.render()
+                with gr.Row():
+                    token_count.render()
+                    char_count.render()
+
+        with gr.Column():
+            output_html.render()
 
     # 定义CSS样式
     webui.css = """
@@ -137,8 +179,15 @@ with gr.Blocks(title="Token Visualizer", theme="NoCrypt/miku") as webui:
 
     submit_btn.click(
         fn=tokenize_text,
-        inputs=[model_selector, custom_model, text_input],
-        outputs=[output_html, token_count, char_count],
+        inputs=[model_selector, text_input],
+        outputs=[
+            output_html,
+            tokenizer_type,
+            vocab_size,
+            sp_tokens,
+            token_count,
+            char_count,
+        ],
     )
 
 if __name__ == "__main__":