feat: support WizardLM (#75)

Added support for building WizardLM

Author: idosal

build.py

@@ -78,7 +78,7 @@ def debug_dump_shader(ex, name, args):
 
 
 def get_models(config, model):
-    if "vicuna" in model or "llama" in model:
+    if "vicuna" in model or "llama" in model or "wizardlm" in model:
         bb = relax.BlockBuilder()
         llama.create_encoding_func(bb, config)
         llama.create_decoding_func(bb, config)

scripts/build_site.sh

@@ -24,5 +24,10 @@ if [ -d "dist/vicuna-7b-v1/params" ]; then
     cp -rf dist/models/vicuna-7b-v1/tokenizer.model site/dist/vicuna-7b-v1/
     cp -rf dist/vicuna-7b-v1/vicuna-7b-v1_webgpu.wasm site/dist/vicuna-7b-v1/
 fi
+if [ -d "dist/wizardlm-7b/params" ]; then
+    mkdir -p site/dist/wizardlm-7b
+    cp -rf dist/models/wizardlm-7b/tokenizer.model site/dist/wizardlm-7b/
+    cp -rf dist/wizardlm-7b/wizardlm-7b_webgpu.wasm site/dist/wizardlm-7b/
+fi
 
 cd site && jekyll b && cd ..

scripts/local_deploy_site.sh

@@ -9,4 +9,9 @@ if [ -d "dist/vicuna-7b-v1/params" ]; then
     rm -rf site/_site/vicuna-7b-v1-params
     ln -s `pwd`/dist/vicuna-7b-v1/params site/_site/vicuna-7b-v1-params
 fi
+if [ -d "dist/wizardlm-7b/params" ]; then
+    rm -rf site/_site/wizardlm-7b-params
+    ln -s `pwd`/dist/wizardlm-7b/params site/_site/wizardlm-7b-params
+fi
+
 cd site && jekyll serve  --skip-initial-build --host localhost --baseurl /web-llm --port 8888

scripts/rpc_debug_deploy.sh

@@ -26,3 +26,11 @@ if [ -d "dist/vicuna-7b-v1/params" ]; then
     rm -rf ${TVM_HOME}/web/.ndarray_cache/vicuna-7b-v1-params
     ln -s `pwd`/dist/vicuna-7b-v1/params ${TVM_HOME}/web/.ndarray_cache/vicuna-7b-v1-params
 fi
+
+if [ -d "dist/wizardlm-7b/params" ]; then
+    mkdir -p ${TVM_HOME}/web/dist/www/dist/wizardlm-7b
+    cp -rf dist/models/wizardlm-7b/tokenizer.model ${TVM_HOME}/web/dist/www/dist/wizardlm-7b/
+    cp -rf dist/wizardlm-7b/wizardlm-7b_webgpu.wasm ${TVM_HOME}/web/dist/www/dist/wizardlm-7b/
+    rm -rf ${TVM_HOME}/web/.ndarray_cache/wizardlm-7b-params
+    ln -s `pwd`/dist/wizardlm-7b/params ${TVM_HOME}/web/.ndarray_cache/wizardlm-7b-params
+fi

web_llm/conversation.py

@@ -160,6 +160,15 @@ def dict(self):
     sep2="</s>",
 )
 
+conv_wizardlm = Conversation(
+    system="You are an AI assistant that gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("", "### Response"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep="\n\n",
+    sep2="</s>",
+)
 
 conv_koala_v1 = Conversation(
     system="BEGINNING OF CONVERSATION:",
@@ -187,6 +196,7 @@ def dict(self):
     "vicuna_v1.1": conv_vicuna_v1_1,
     "koala_v1": conv_koala_v1,
     "dolly": conv_dolly,
+    "wizardlm": conv_wizardlm
 }
 
 
@@ -198,4 +208,6 @@ def get_default_conv_template(model_name):
         return conv_koala_v1
     elif "dolly" in model_name:
         return conv_dolly
+    elif "wizardlm" in model_name:
+        return conv_wizardlm
     return conv_one_shot

web_llm/utils.py

@@ -6,7 +6,7 @@
 
 
 def get_config(hf_config, model):
-    if "vicuna" in model or "llama" in model:
+    if "vicuna" in model or "llama" in model or "wizardlm" in model:
         from .relax_model.llama import LlamaConfig as RelaxConfig
 
         return RelaxConfig(