sglang/docs/embedding_model.ipynb

Embedding Model

SGLang supports embedding models in the same way as completion models. Here are some example models:

• intfloat/e5-mistral-7b-instruct
• Alibaba-NLP/gte-Qwen2-7B-instruct

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 --log-level error

Remember to add --is-embedding to the command.

In [7]:

from sglang.utils import execute_shell_command, wait_for_server, terminate_process

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 --log-level error
"""
)

wait_for_server("http://localhost:30010")

print("Embedding server is ready. Proceeding with the next steps.")

Embedding server is ready. Proceeding with the next steps.

Use Curl

In [8]:

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(f"Text embedding (first 10): {text_embedding[:10]}")

Text embedding (first 10): [0.0083160400390625, 0.0006804466247558594, -0.00809478759765625, -0.0006995201110839844, 0.0143890380859375, -0.0090179443359375, 0.01238250732421875, 0.00209808349609375, 0.0062103271484375, -0.003047943115234375]

Using OpenAI Compatible API

In [9]:

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(f"Text embedding (first 10): {embedding}")

Text embedding (first 10): [0.00829315185546875, 0.0007004737854003906, -0.00809478759765625, -0.0006799697875976562, 0.01438140869140625, -0.00897979736328125, 0.0123748779296875, 0.0020923614501953125, 0.006195068359375, -0.0030498504638671875]

Using Input IDs

SGLang also supports input_ids as input to get the embedding.

In [10]:

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(f"Input IDs embedding (first 10): {input_ids_embedding[:10]}")

Input IDs embedding (first 10): [0.00829315185546875, 0.0007004737854003906, -0.00809478759765625, -0.0006799697875976562, 0.01438140869140625, -0.00897979736328125, 0.0123748779296875, 0.0020923614501953125, 0.006195068359375, -0.0030498504638671875]

In [11]:

terminate_process(embedding_process)