diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..567bff5 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,25 @@ +FROM registry.maas.sunrise-ai.com/public/vllm:S2-v1.1.1 +ENV LD_LIBRARY_PATH=/usr/local/pccl/lib:\ +/usr/local/tangrt/targets/linux-x86_64/lib:\ +/usr/local/tangrt/targets/linux-x86_64/lib/stub:\ +/root/pt200/gcc-11.3.0/install/lib64:\ +/root:/root/gcc-11.5.0/lib64:\ +/usr/local/pccl/lib:\ +/usr/local/tangrt/targets/linux-x86_64/lib:\ +/usr/local/tangrt/targets/linux-x86_64/lib/stub:\ +/usr/local/tangrt/lib/linux-x86_64:\ +/root/pt200/gcc-11.3.0/install/lib64:\ +/root:\ +/usr/lib64:\ +/usr/local/lib/python3.10/site-packages/torch/lib +ENV TORCH_DEVICE_BACKEND_AUTOLOAD=0 +ENV PATH=/root/gcc-11.5.0/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENV PYTHONPATH=/sunrise_code/vllm:/sunrise_code/sunrise_vllm:/usr/local/lib/python3.10/site-packages: +RUN ln -sf /usr/local/bin/python3.10 /usr/bin/python3 +COPY fix_tokenizer.py /opt/ +COPY detect_tokenizer.py /opt/ +COPY entrypoint.sh /opt/ +WORKDIR /model +COPY . /model +RUN chmod +x /opt/entrypoint.sh +ENTRYPOINT ["/opt/entrypoint.sh"] diff --git a/detect_tokenizer.py b/detect_tokenizer.py new file mode 100644 index 0000000..9a87f80 --- /dev/null +++ b/detect_tokenizer.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 +# 检测 tokenizer 是否需要修复 +import json +import os + +MODEL_DIR = os.environ.get("MODEL_DIR", "/model") +cfg_path = os.path.join(MODEL_DIR, "tokenizer_config.json") +if not os.path.exists(cfg_path): + print("no tokenizer_config.json") + raise SystemExit(0) +with open(cfg_path, "r", encoding="utf-8") as f: + cfg = json.load(f) +if cfg.get("tokenizer_class") in {"TokenizersBackend", "TiktokenTokenizer"}: + print("need_fix") +elif isinstance(cfg.get("extra_special_tokens"), list): + print("need_fix") +else: + print("ok") diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100644 index 0000000..a8ad73a --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,38 @@ +#!/bin/bash +set -e + +MODEL_DIR=${1:-/model} +shift || true + +FIX_TOKENIZER_DIR=/tmp/fixed_tokenizer +AUTO_FIX=${AUTO_FIX_TOKENIZER:-auto} + +echo "[entrypoint] model dir: $MODEL_DIR" + +NEED_FIX=0 + +if [ "$AUTO_FIX" = "1" ] || [ "$AUTO_FIX" = "true" ]; then + NEED_FIX=1 +elif [ "$AUTO_FIX" = "auto" ]; then + if [ -f "$MODEL_DIR/tokenizer_config.json" ]; then + if grep -q "TokenizersBackend\|TiktokenTokenizer" "$MODEL_DIR/tokenizer_config.json"; then + NEED_FIX=1 + fi + if grep -q '"extra_special_tokens":\s*\[' "$MODEL_DIR/tokenizer_config.json"; then + NEED_FIX=1 + fi + fi +fi + +if [ $NEED_FIX -eq 1 ]; then + echo "[entrypoint] fixing tokenizer..." + python3 /opt/fix_tokenizer.py + TOKENIZER_ARG="--tokenizer $FIX_TOKENIZER_DIR" +else + echo "[entrypoint] tokenizer OK, skip fix" + TOKENIZER_ARG="" +fi + +echo "[entrypoint] starting vllm..." + +exec vllm serve "$MODEL_DIR" $TOKENIZER_ARG "$@" diff --git a/fix_tokenizer.py b/fix_tokenizer.py new file mode 100644 index 0000000..70092c4 --- /dev/null +++ b/fix_tokenizer.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +# 自动修复 tokenizer 配置,供 entrypoint.sh 调用 +import json +import os +import shutil + +MODEL_DIR = os.environ.get("MODEL_DIR", "/model") +FIX_DIR = "/tmp/fixed_tokenizer" + +os.makedirs(FIX_DIR, exist_ok=True) +for name in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.json", "merges.txt", "tokenizer.model"]: + src = os.path.join(MODEL_DIR, name) + if os.path.exists(src): + shutil.copy(src, os.path.join(FIX_DIR, name)) + +cfg_path = os.path.join(FIX_DIR, "tokenizer_config.json") +if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding="utf-8") as f: + cfg = json.load(f) + changed = False + bad = {"TokenizersBackend", "TiktokenTokenizer"} + if cfg.get("tokenizer_class") in bad: + cfg["tokenizer_class"] = "PreTrainedTokenizerFast" + changed = True + extra = cfg.get("extra_special_tokens") + if isinstance(extra, list): + cfg["extra_special_tokens"] = {t: t for t in extra} + changed = True + if changed: + with open(cfg_path, "w", encoding="utf-8") as f: + json.dump(cfg, f, ensure_ascii=False, indent=2) + print("[fix] tokenizer_config.json 已修复")