Embedding Model#
SGLang supports embedding models in the same way as completion models. Here are some example models:
Launch A Server#
The following code is equivalent to running this in the shell:
python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct \
--port 30010 --host 0.0.0.0 --is-embedding
Remember to add --is-embedding to the command.
[1]:
from sglang.utils import (
execute_shell_command,
wait_for_server,
terminate_process,
print_highlight,
)
embedding_process = execute_shell_command(
"""
python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct \
--port 30010 --host 0.0.0.0 --is-embedding
"""
)
wait_for_server("http://localhost:30010")
[2024-11-02 03:44:18] server_args=ServerArgs(model_path='Alibaba-NLP/gte-Qwen2-7B-instruct', tokenizer_path='Alibaba-NLP/gte-Qwen2-7B-instruct', tokenizer_mode='auto', skip_tokenizer_init=False, load_format='auto', trust_remote_code=False, dtype='auto', kv_cache_dtype='auto', quantization=None, context_length=None, device='cuda', served_model_name='Alibaba-NLP/gte-Qwen2-7B-instruct', chat_template=None, is_embedding=True, host='0.0.0.0', port=30010, mem_fraction_static=0.88, max_running_requests=None, max_total_tokens=None, chunked_prefill_size=8192, max_prefill_tokens=16384, schedule_policy='lpm', schedule_conservativeness=1.0, tp_size=1, stream_interval=1, random_seed=72474417, constrained_json_whitespace_pattern=None, decode_log_interval=40, log_level='info', log_level_http=None, log_requests=False, show_time_cost=False, api_key=None, file_storage_pth='SGLang_storage', enable_cache_report=False, watchdog_timeout=600, dp_size=1, load_balance_method='round_robin', dist_init_addr=None, nnodes=1, node_rank=0, json_model_override_args='{}', enable_double_sparsity=False, ds_channel_config_path=None, ds_heavy_channel_num=32, ds_heavy_token_num=256, ds_heavy_channel_type='qk', ds_sparse_decode_threshold=4096, lora_paths=None, max_loras_per_batch=8, attention_backend='flashinfer', sampling_backend='flashinfer', grammar_backend='outlines', disable_flashinfer=False, disable_flashinfer_sampling=False, disable_radix_cache=False, disable_regex_jump_forward=False, disable_cuda_graph=False, disable_cuda_graph_padding=False, disable_disk_cache=False, disable_custom_all_reduce=False, disable_mla=False, disable_penalizer=False, disable_nan_detection=False, enable_overlap_schedule=False, enable_mixed_chunk=False, enable_torch_compile=False, torch_compile_max_bs=32, cuda_graph_max_bs=160, torchao_config='', enable_p2p_check=False, triton_attention_reduce_in_fp32=False, num_continuous_decode_steps=1)
[2024-11-02 03:44:33 TP0] Init torch distributed begin.
[2024-11-02 03:44:34 TP0] Load weight begin. avail mem=78.59 GB
[2024-11-02 03:44:34 TP0] lm_eval is not installed, GPTQ may not be usable
INFO 11-02 03:44:34 weight_utils.py:243] Using model weights format ['*.safetensors']
Loading safetensors checkpoint shards: 0% Completed | 0/7 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 14% Completed | 1/7 [00:01<00:10, 1.83s/it]
Loading safetensors checkpoint shards: 29% Completed | 2/7 [00:03<00:09, 1.87s/it]
Loading safetensors checkpoint shards: 43% Completed | 3/7 [00:05<00:07, 1.93s/it]
Loading safetensors checkpoint shards: 57% Completed | 4/7 [00:07<00:05, 1.95s/it]
Loading safetensors checkpoint shards: 71% Completed | 5/7 [00:09<00:03, 1.96s/it]
Loading safetensors checkpoint shards: 86% Completed | 6/7 [00:11<00:01, 1.80s/it]
Loading safetensors checkpoint shards: 100% Completed | 7/7 [00:12<00:00, 1.53s/it]
Loading safetensors checkpoint shards: 100% Completed | 7/7 [00:12<00:00, 1.74s/it]
[2024-11-02 03:44:47 TP0] Load weight end. type=Qwen2ForCausalLM, dtype=torch.float16, avail mem=64.18 GB
[2024-11-02 03:44:47 TP0] Memory pool end. avail mem=7.43 GB
[2024-11-02 03:44:47 TP0] max_total_num_tokens=1025173, max_prefill_tokens=16384, max_running_requests=4005, context_len=131072
[2024-11-02 03:44:47] INFO: Started server process [3453536]
[2024-11-02 03:44:47] INFO: Waiting for application startup.
[2024-11-02 03:44:47] INFO: Application startup complete.
[2024-11-02 03:44:47] INFO: Uvicorn running on http://0.0.0.0:30010 (Press CTRL+C to quit)
[2024-11-02 03:44:48] INFO: 127.0.0.1:55542 - "GET /v1/models HTTP/1.1" 200 OK
[2024-11-02 03:44:48] INFO: 127.0.0.1:55562 - "GET /get_model_info HTTP/1.1" 200 OK
[2024-11-02 03:44:48 TP0] Prefill batch. #new-seq: 1, #new-token: 6, #cached-token: 0, cache hit rate: 0.00%, token usage: 0.00, #running-req: 0, #queue-req: 0
[2024-11-02 03:44:50] INFO: 127.0.0.1:55572 - "POST /encode HTTP/1.1" 200 OK
[2024-11-02 03:44:50] The server is fired up and ready to roll!
NOTE: Typically, the server runs in a separate terminal.
In this notebook, we run the server and notebook code together, so their outputs are combined.
To improve clarity, the server logs are displayed in the original black color, while the notebook outputs are highlighted in blue.
Use Curl#
[2]:
import subprocess, json
text = "Once upon a time"
curl_text = f"""curl -s http://localhost:30010/v1/embeddings \
-H "Content-Type: application/json" \
-H "Authorization: Bearer None" \
-d '{{"model": "Alibaba-NLP/gte-Qwen2-7B-instruct", "input": "{text}"}}'"""
text_embedding = json.loads(subprocess.check_output(curl_text, shell=True))["data"][0][
"embedding"
]
print_highlight(f"Text embedding (first 10): {text_embedding[:10]}")
[2024-11-02 03:44:53 TP0] Prefill batch. #new-seq: 1, #new-token: 4, #cached-token: 0, cache hit rate: 0.00%, token usage: 0.00, #running-req: 0, #queue-req: 0
[2024-11-02 03:44:53] INFO: 127.0.0.1:55634 - "POST /v1/embeddings HTTP/1.1" 200 OK
Text embedding (first 10): [0.00830841064453125, 0.0006804466247558594, -0.00807952880859375, -0.000682830810546875, 0.01438140869140625, -0.009002685546875, 0.01239013671875, 0.0020999908447265625, 0.006214141845703125, -0.0030345916748046875]
Using OpenAI Compatible API#
[3]:
import openai
client = openai.Client(base_url="http://127.0.0.1:30010/v1", api_key="None")
# Text embedding example
response = client.embeddings.create(
model="Alibaba-NLP/gte-Qwen2-7B-instruct",
input=text,
)
embedding = response.data[0].embedding[:10]
print_highlight(f"Text embedding (first 10): {embedding}")
[2024-11-02 03:44:53 TP0] Prefill batch. #new-seq: 1, #new-token: 1, #cached-token: 3, cache hit rate: 21.43%, token usage: 0.00, #running-req: 0, #queue-req: 0
[2024-11-02 03:44:53] INFO: 127.0.0.1:55650 - "POST /v1/embeddings HTTP/1.1" 200 OK
Text embedding (first 10): [0.00830078125, 0.0006747245788574219, -0.00807952880859375, -0.000682830810546875, 0.01438140869140625, -0.009002685546875, 0.01239013671875, 0.0020961761474609375, 0.006198883056640625, -0.003025054931640625]
Using Input IDs#
SGLang also supports input_ids as input to get the embedding.
[4]:
import json
import os
from transformers import AutoTokenizer
os.environ["TOKENIZERS_PARALLELISM"] = "false"
tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-Qwen2-7B-instruct")
input_ids = tokenizer.encode(text)
curl_ids = f"""curl -s http://localhost:30010/v1/embeddings \
-H "Content-Type: application/json" \
-H "Authorization: Bearer None" \
-d '{{"model": "Alibaba-NLP/gte-Qwen2-7B-instruct", "input": {json.dumps(input_ids)}}}'"""
input_ids_embedding = json.loads(subprocess.check_output(curl_ids, shell=True))["data"][
0
]["embedding"]
print_highlight(f"Input IDs embedding (first 10): {input_ids_embedding[:10]}")
[2024-11-02 03:45:00 TP0] Prefill batch. #new-seq: 1, #new-token: 1, #cached-token: 3, cache hit rate: 33.33%, token usage: 0.00, #running-req: 0, #queue-req: 0
[2024-11-02 03:45:00] INFO: 127.0.0.1:52916 - "POST /v1/embeddings HTTP/1.1" 200 OK
Input IDs embedding (first 10): [0.00830078125, 0.0006747245788574219, -0.00807952880859375, -0.000682830810546875, 0.01438140869140625, -0.009002685546875, 0.01239013671875, 0.0020961761474609375, 0.006198883056640625, -0.003025054931640625]
[5]:
terminate_process(embedding_process)