Update app.py

app.py

@@ -4,7 +4,7 @@ from html import escape
 from transformers import AutoTokenizer
 
 
-def get_available_models():
+def get_available_models() -> list[str]:
     """获取models目录下所有包含tokenizer.json的模型"""
     models_dir = "models"
     if not os.path.exists(models_dir):
@@ -13,7 +13,7 @@ def get_available_models():
     available_models = []
     for model_name in os.listdir(models_dir):
         model_path = os.path.join(models_dir, model_name)
-        tokenizer_file = os.path.join(model_path, "config.json")
+        tokenizer_file = os.path.join(model_path, "tokenizer.json")
 
         if os.path.isdir(model_path) and os.path.isfile(tokenizer_file):
             available_models.append(model_name)
@@ -21,17 +21,26 @@ def get_available_models():
     return sorted(available_models)
 
 
-def tokenize_text(model_name, text):
+def tokenize_text(
+    builtin_model: str, custom_model: str | None, text: str
+) -> tuple[str | None, int, int]:
     """处理tokenize请求"""
-    if not model_name:
+    if not builtin_model:
         return "Please choose a model and input some texts", 0, 0
     if not text:
         text = "Please choose a model and input some texts"
 
     try:
         # 加载tokenizer
-        model_path = os.path.join("models", model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, device_map="cpu")
+        if custom_model:
+            tokenizer = AutoTokenizer.from_pretrained(
+                custom_model, 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"
+            )
 
         # Tokenize处理
         input_ids = tokenizer.encode(text, add_special_tokens=True)
@@ -63,7 +72,7 @@ def tokenize_text(model_name, text):
 
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        return error_msg, ""
+        return error_msg, 0, 0
 
 
 banner_md = """# 🎨 Tokenize it!
@@ -73,7 +82,12 @@ 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="Choose Model", choices=get_available_models(), interactive=True
+    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",
+    lines=1,
 )
 text_input = gr.Textbox(label="Input Text", placeholder="Hello World!", lines=4)
 submit_btn = gr.Button("🚀 Tokenize!", variant="primary")
@@ -86,7 +100,10 @@ with gr.Blocks(title="Token Visualizer", theme="NoCrypt/miku") as webui:
     banner.render()
 
     with gr.Column():
-        model_selector.render()
+        with gr.TabItem("Built-in Model"):
+            model_selector.render()
+        with gr.TabItem("Custom Model"):
+            custom_model.render()
         text_input.render()
         submit_btn.render()
 
@@ -120,7 +137,7 @@ with gr.Blocks(title="Token Visualizer", theme="NoCrypt/miku") as webui:
 
     submit_btn.click(
         fn=tokenize_text,
-        inputs=[model_selector, text_input],
+        inputs=[model_selector, custom_model, text_input],
         outputs=[output_html, token_count, char_count],
     )