Sanitize added_tokens before tokenizer load

app.py

@@ -1,15 +1,39 @@
 import gradio as gr
+import json
+import os
+from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
 # 加载 CodeT5+ 模型
 model_name = "Salesforce/codet5p-220m"
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, trust_remote_code=True)
-except TypeError:
-    # Some tokenizer repos expose added_tokens metadata that breaks fast tokenizer init.
-    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, trust_remote_code=True)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
+
+
+def prepare_local_model(repo_id: str, local_dir: str = "./model_cache") -> str:
+    snapshot_download(repo_id=repo_id, local_dir=local_dir)
+
+    added_tokens_file = os.path.join(local_dir, "added_tokens.json")
+    if os.path.exists(added_tokens_file):
+        with open(added_tokens_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        # Ensure the file is a plain token list for compatibility with tokenizers.add_tokens.
+        if isinstance(data, dict):
+            normalized = list(data.keys())
+        elif isinstance(data, list):
+            normalized = [str(item) for item in data]
+        else:
+            normalized = []
+
+        with open(added_tokens_file, "w", encoding="utf-8") as f:
+            json.dump(normalized, f, ensure_ascii=False)
+
+    return local_dir
+
+
+local_model_dir = prepare_local_model(model_name)
+tokenizer = AutoTokenizer.from_pretrained(local_model_dir, use_fast=True, trust_remote_code=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(local_model_dir, trust_remote_code=True)
 
 def generate_code(prompt: str, max_length: int = 128) -> str:
     """代码生成/补全"""

requirements.txt

@@ -1,4 +1,5 @@
 transformers>=4.40.0
+huggingface_hub>=0.23.0
 torch>=2.0.0
 sentencepiece>=0.1.96
 accelerate>=0.20.0