59799a3687
Co-authored-by: zhouhaizhu.zhz <zhouhaizhu.zhz@alibaba-inc.com> Co-authored-by: LinSiyuan814 <linsiyuan.lsy@alibaba-inc.com> Co-authored-by: hzh0425 <hzh0425@apache.org>
56 lines
1.4 KiB
Python
Executable File
56 lines
1.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Rebuild the small ML and tokenizer assets used by examples and tests."""
|
|
|
|
from pathlib import Path
|
|
|
|
import joblib
|
|
import numpy as np
|
|
from sglang_simulator.time_predictor.ml import MLTimePredictor
|
|
from sklearn.dummy import DummyRegressor
|
|
from tokenizers import Tokenizer
|
|
from tokenizers.models import WordLevel
|
|
from tokenizers.pre_tokenizers import Whitespace
|
|
from transformers import PreTrainedTokenizerFast
|
|
|
|
ASSETS = Path(__file__).parent / "assets"
|
|
|
|
|
|
def build_ml_model() -> None:
|
|
model = DummyRegressor(strategy="constant", constant=0.001)
|
|
model.fit(np.zeros((1, len(MLTimePredictor.FEATURE_NAMES))), [0.001])
|
|
joblib.dump(
|
|
{"model": model, "features": MLTimePredictor.FEATURE_NAMES},
|
|
ASSETS / "model.pkl",
|
|
)
|
|
|
|
|
|
def build_tokenizer() -> None:
|
|
tokenizer = Tokenizer(
|
|
WordLevel(
|
|
{
|
|
"[UNK]": 0,
|
|
"prefix": 1,
|
|
"caching": 2,
|
|
"latency": 3,
|
|
"decode": 4,
|
|
"token": 5,
|
|
},
|
|
unk_token="[UNK]",
|
|
)
|
|
)
|
|
tokenizer.pre_tokenizer = Whitespace()
|
|
PreTrainedTokenizerFast(
|
|
tokenizer_object=tokenizer,
|
|
unk_token="[UNK]",
|
|
).save_pretrained(ASSETS / "tokenizer")
|
|
|
|
|
|
def main() -> None:
|
|
ASSETS.mkdir(parents=True, exist_ok=True)
|
|
build_ml_model()
|
|
build_tokenizer()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|