From 9fc3e8aac7422826a7fad3477ddafe74e03e6420 Mon Sep 17 00:00:00 2001 From: satyamk7054 <43010011+satyamk7054@users.noreply.github.com> Date: Mon, 27 Oct 2025 11:49:36 -0700 Subject: [PATCH] Add support for Matryoshka embeddings (#126) (#11142) Co-authored-by: Satyam Kumar --- benchmark/prefill_only/bench_embeddings.py | 15 ++- docs/supported_models/embedding_models.md | 39 +++++++ python/sglang/srt/configs/model_config.py | 8 ++ python/sglang/srt/entrypoints/engine.py | 4 + .../entrypoints/openai/serving_embedding.py | 1 + python/sglang/srt/layers/pooler.py | 23 +++- python/sglang/srt/managers/io_struct.py | 6 ++ python/sglang/srt/managers/schedule_batch.py | 20 ++++ python/sglang/srt/managers/scheduler.py | 1 + .../scheduler_output_processor_mixin.py | 5 +- .../sglang/srt/managers/tokenizer_manager.py | 33 ++++++ .../srt/model_executor/forward_batch_info.py | 4 + python/sglang/test/runners.py | 36 +++++-- test/srt/models/test_embedding_models.py | 29 ++++- .../basic/test_openai_embedding.py | 102 ++++++++++++++++++ 15 files changed, 314 insertions(+), 12 deletions(-) diff --git a/benchmark/prefill_only/bench_embeddings.py b/benchmark/prefill_only/bench_embeddings.py index ca66c85a3..74d8a582e 100644 --- a/benchmark/prefill_only/bench_embeddings.py +++ b/benchmark/prefill_only/bench_embeddings.py @@ -18,6 +18,7 @@ Usage: import asyncio import logging +from typing import Optional from transformers import AutoTokenizer from util import ( @@ -52,11 +53,14 @@ config.freeze_gc = True # Enable GC freeze functionality HTTP_URL = "http://localhost:30000/v1/embeddings" # Embeddings API Config -EMBEDDINGS_MODEL_PATH = "/Qwen/Qwen3-Embedding-0.6B" +EMBEDDINGS_MODEL_PATH = "Qwen/Qwen3-Embedding-0.6B" BATCH_SIZE = [1] # Number of items per request (batch size) # Configurable input token length EMBEDDINGS_INPUT_TOKENS = 500 # Default token length +MATRYOSHKA_DIMENSIONS: Optional[int] = ( + None # Set to None to disable matryoshka embeddings +) # Load tokenizer once for embeddings text generation print("Loading tokenizer for embeddings input generation...") @@ -85,6 +89,7 @@ def build_embeddings_request(index: int, item_count: int) -> tuple: req = { "input": input_data, "model": EMBEDDINGS_MODEL_PATH, + "dimensions": MATRYOSHKA_DIMENSIONS, } return (index, req) except Exception as e: @@ -94,7 +99,12 @@ def build_embeddings_request(index: int, item_count: int) -> tuple: def validate_embeddings_response(response_data: dict) -> bool: """Validate embeddings API response.""" - return "data" in response_data + return ( + "data" in response_data + and len(response_data["data"][0]["embedding"]) == MATRYOSHKA_DIMENSIONS + if MATRYOSHKA_DIMENSIONS + else True + ) def build_warmup_embeddings_request() -> dict: @@ -102,6 +112,7 @@ def build_warmup_embeddings_request() -> dict: return { "input": EMBEDDINGS_INPUT_TEXT, "model": EMBEDDINGS_MODEL_PATH, + "dimensions": MATRYOSHKA_DIMENSIONS, } diff --git a/docs/supported_models/embedding_models.md b/docs/supported_models/embedding_models.md index 437cb8284..906466ac5 100644 --- a/docs/supported_models/embedding_models.md +++ b/docs/supported_models/embedding_models.md @@ -75,6 +75,45 @@ response = requests.post(url + "/v1/embeddings", json=payload).json() print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) ``` +## Matryoshka Embedding Example + +[Matryoshka Embeddings](https://sbert.net/examples/sentence_transformer/training/matryoshka/README.html#matryoshka-embeddings) or [Matryoshka Representation Learning (MRL)](https://arxiv.org/abs/2205.13147) is a technique used in training embedding models. It allows user to trade off between performance and cost. + +### 1. Launch a Matryoshka‑capable model + +If the model config already includes `matryoshka_dimensions` or `is_matryoshka` then no override is needed. Otherwise, you can use `--json-model-override-args` as below: + +```shell +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-Embedding-0.6B \ + --is-embedding \ + --host 0.0.0.0 \ + --port 30000 \ + --json-model-override-args '{"matryoshka_dimensions": [128, 256, 512, 1024, 1536]}' +``` + +1. Setting `"is_matryoshka": true` allows truncating to any dimension. Otherwise, the server will validate that the specified dimension in the request is one of `matryoshka_dimensions`. +2. Omitting `dimensions` in a request returns the full vector. + +### 2. Make requests with different output dimensions + +```python +import requests + +url = "http://127.0.0.1:30000" + +# Request a truncated (Matryoshka) embedding by specifying a supported dimension. +payload = { + "model": "Qwen/Qwen3-Embedding-0.6B", + "input": "Explain diffusion models simply.", + "dimensions": 512 # change to 128 / 1024 / omit for full size +} + +response = requests.post(url + "/v1/embeddings", json=payload).json() +print("Embedding:", response["data"][0]["embedding"]) +``` + + ## Supported Models | Model Family | Example Model | Chat Template | Description | diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 508fe9bcf..69a4a545c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -205,6 +205,14 @@ class ModelConfig: self.hf_config, "image_token_id", None ) or getattr(self.hf_config, "image_token_index", None) + # matryoshka embeddings + self.matryoshka_dimensions = getattr( + self.hf_config, "matryoshka_dimensions", None + ) + self.is_matryoshka = self.matryoshka_dimensions or getattr( + self.hf_config, "is_matryoshka", False + ) + @staticmethod def from_server_args( server_args: ServerArgs, diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index f79a55143..7b0ebb79f 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -312,6 +312,7 @@ class Engine(EngineBase): image_data: Optional[MultimodalDataInputFormat] = None, audio_data: Optional[MultimodalDataInputFormat] = None, video_data: Optional[MultimodalDataInputFormat] = None, + dimensions: Optional[int] = None, ) -> Dict: """ The arguments of this function is the same as `sglang/srt/managers/io_struct.py::EmbeddingReqInput`. @@ -322,6 +323,7 @@ class Engine(EngineBase): image_data=image_data, audio_data=audio_data, video_data=video_data, + dimensions=dimensions, ) generator = self.tokenizer_manager.generate_request(obj, None) ret = self.loop.run_until_complete(generator.__anext__()) @@ -333,6 +335,7 @@ class Engine(EngineBase): image_data: Optional[MultimodalDataInputFormat] = None, audio_data: Optional[MultimodalDataInputFormat] = None, video_data: Optional[MultimodalDataInputFormat] = None, + dimensions: Optional[int] = None, ) -> Dict: """ Asynchronous version of encode method. @@ -345,6 +348,7 @@ class Engine(EngineBase): image_data=image_data, audio_data=audio_data, video_data=video_data, + dimensions=dimensions, ) generator = self.tokenizer_manager.generate_request(obj, None) return await generator.__anext__() diff --git a/python/sglang/srt/entrypoints/openai/serving_embedding.py b/python/sglang/srt/entrypoints/openai/serving_embedding.py index 7340a72f2..08e48ddd4 100644 --- a/python/sglang/srt/entrypoints/openai/serving_embedding.py +++ b/python/sglang/srt/entrypoints/openai/serving_embedding.py @@ -126,6 +126,7 @@ class OpenAIServingEmbedding(OpenAIServingBase): **prompt_kwargs, rid=request.rid, priority=request.priority, + dimensions=request.dimensions, ) return adapted_request, request diff --git a/python/sglang/srt/layers/pooler.py b/python/sglang/srt/layers/pooler.py index 26bc5899e..2c98c856b 100644 --- a/python/sglang/srt/layers/pooler.py +++ b/python/sglang/srt/layers/pooler.py @@ -20,7 +20,9 @@ class PoolingType(IntEnum): @dataclass class EmbeddingPoolerOutput: - embeddings: torch.Tensor + # Pooler can return list[tensor] instead of tensor if the dimension of each tensor in the batch is different + # due to different per-request matryoshka dim truncation + embeddings: torch.Tensor | list[torch.Tensor] class Pooler(nn.Module): @@ -42,6 +44,7 @@ class Pooler(nn.Module): def forward( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ) -> EmbeddingPoolerOutput: + if self.pooling_type == PoolingType.LAST: last_token_indices = torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1 pooled_data = hidden_states[last_token_indices] @@ -53,8 +56,24 @@ class Pooler(nn.Module): else: raise ValueError(f"Invalid pooling type: {self.pooling_type}") + if forward_batch.dimensions is not None: + all_same_dimensions = len(set(forward_batch.dimensions)) == 1 + if all_same_dimensions: + pooled_data = pooled_data[..., : forward_batch.dimensions[0]] + else: + pooled_data = [ + tensor[..., :dim] + for tensor, dim in zip(pooled_data, forward_batch.dimensions) + ] + if self.normalize: - pooled_data = nn.functional.normalize(pooled_data, p=2, dim=1) + if isinstance(pooled_data, list): + pooled_data = [ + nn.functional.normalize(tensor, p=2, dim=-1) + for tensor in pooled_data + ] + else: + pooled_data = nn.functional.normalize(pooled_data, p=2, dim=-1) return EmbeddingPoolerOutput(embeddings=pooled_data) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index d7671e3e8..78b089c42 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -695,6 +695,9 @@ class EmbeddingReqInput(BaseReq): # tracing context trace_context: Optional[Dict] = None + # The number of dimensions the resulting output embeddings should have. It is applicable for Matryoshka Embeddings. + dimensions: Optional[int] = None + def normalize_batch_and_arguments(self): # at least one of text, input_ids, or image should be provided if self.text is None and self.input_ids is None and self.image_data is None: @@ -771,6 +774,7 @@ class EmbeddingReqInput(BaseReq): video_data=self.video_data[i] if self.video_data is not None else None, sampling_params=self.sampling_params[i], rid=self.rid[i], + dimensions=self.dimensions, http_worker_ipc=self.http_worker_ipc, ) @@ -791,6 +795,8 @@ class TokenizedEmbeddingReqInput(BaseReq): data_parallel_rank: Optional[int] = None # Priority for the request priority: Optional[int] = None + # The number of dimensions the resulting output embeddings should have. It is applicable for Matryoshka Embeddings. + dimensions: Optional[int] = None @dataclass diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3d625ac9e..953a4c62f 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -442,6 +442,7 @@ class Req: priority: Optional[int] = None, metrics_collector: Optional[SchedulerMetricsCollector] = None, extra_key: Optional[str] = None, + dimensions: Optional[int] = None, http_worker_ipc: Optional[str] = None, ): # Input and output info @@ -650,6 +651,9 @@ class Req: self.tmp_end_idx: int = -1 self.metadata_buffer_index: int = -1 + # For Matryoshka embeddings + self.dimensions = dimensions + @property def seqlen(self): return len(self.origin_input_ids) + len(self.output_ids) @@ -1014,6 +1018,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): encoder_lens_cpu: Optional[List[int]] = None encoder_out_cache_loc: Optional[torch.Tensor] = None + # For matryoshka embeddings + dimensions: Optional[list[int]] = None + # For split prefill split_index: int = 0 split_prefill_finished: bool = False @@ -1177,6 +1184,15 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): prefix_lens = [len(r.prefix_indices) for r in reqs] extend_lens = [r.extend_input_len for r in reqs] + # For matryoshka embeddings + if self.model_config.is_matryoshka and any( + r.dimensions is not None for r in reqs + ): + self.dimensions = [ + r.dimensions if r.dimensions else self.model_config.hidden_size + for r in reqs + ] + token_type_ids = [ r.token_type_ids for r in reqs if r.token_type_ids is not None ] @@ -1765,6 +1781,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ), extend_input_logprob_token_ids=self.extend_input_logprob_token_ids, is_prefill_only=self.is_prefill_only, + dimensions=self.dimensions, ) def copy(self): @@ -1873,5 +1890,8 @@ class ModelWorkerBatch: capture_hidden_mode: CaptureHiddenMode = None hicache_consumer_index: int = -1 + # For matryoshka embeddings + dimensions: Optional[list[int]] = None + # Whether this batch is prefill-only (no token generation needed) is_prefill_only: bool = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 76f2e2f2d..0976b8ec5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1475,6 +1475,7 @@ class Scheduler( recv_req.sampling_params, token_type_ids=recv_req.token_type_ids, priority=recv_req.priority, + dimensions=recv_req.dimensions, http_worker_ipc=recv_req.http_worker_ipc, ) req.tokenizer = self.tokenizer diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index e06fac95a..8f9ac10e9 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -203,7 +203,10 @@ class SchedulerOutputProcessorMixin: i ].item() else: - embeddings = embeddings.tolist() + if isinstance(embeddings, torch.Tensor): + embeddings = embeddings.tolist() + else: + embeddings = [tensor.tolist() for tensor in embeddings] # Check finish conditions for i, req in enumerate(batch.reqs): diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 885da6a98..ba1c914df 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -666,6 +666,10 @@ class TokenizerManager(TokenizerCommunicatorMixin): ) raise ValueError(error_msg) + # Matryoshka embeddings validations + if isinstance(obj, EmbeddingReqInput): + self._validate_for_matryoshka_dim(obj) + if isinstance(obj, GenerateReqInput): if ( obj.return_hidden_states @@ -684,6 +688,34 @@ class TokenizerManager(TokenizerCommunicatorMixin): "Please set `--enable-custom-logit-processor` to enable this feature." ) + def _validate_for_matryoshka_dim(self, obj: EmbeddingReqInput) -> None: + """Validate the request for Matryoshka dim if it has the field set.""" + if obj.dimensions is None: + return + + if not self.model_config.is_matryoshka: + raise ValueError( + f"Model '{self.model_config.model_path}' does not support matryoshka representation, " + f"changing output dimensions will lead to poor results." + ) + + if obj.dimensions < 1: + raise ValueError("Requested dimensions must be greater than 0") + + if ( + self.model_config.matryoshka_dimensions + and obj.dimensions not in self.model_config.matryoshka_dimensions + ): + raise ValueError( + f"Model '{self.model_config.model_path}' only supports {self.model_config.matryoshka_dimensions} matryoshka dimensions, " + f"using other output dimensions will lead to poor results." + ) + + if obj.dimensions > self.model_config.hidden_size: + raise ValueError( + f"Provided dimensions are greater than max embedding dimension: {self.model_config.hidden_size}" + ) + def _validate_input_ids_in_vocab( self, input_ids: List[int], vocab_size: int ) -> None: @@ -752,6 +784,7 @@ class TokenizerManager(TokenizerCommunicatorMixin): sampling_params, rid=obj.rid, priority=obj.priority, + dimensions=obj.dimensions, http_worker_ipc=obj.http_worker_ipc, ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index cbda58cf9..7afcf6fa9 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -320,6 +320,9 @@ class ForwardBatch: tbo_parent_token_range: Optional[Tuple[int, int]] = None tbo_children: Optional[List[ForwardBatch]] = None + # For matryoshka embeddings + dimensions: Optional[list[int]] = None + @classmethod def init_new( cls, @@ -361,6 +364,7 @@ class ForwardBatch: input_embeds=batch.input_embeds, token_type_ids=batch.token_type_ids, tbo_split_seq_index=batch.tbo_split_seq_index, + dimensions=batch.dimensions, ) device = model_runner.device diff --git a/python/sglang/test/runners.py b/python/sglang/test/runners.py index dc7efe528..76cfbcb0b 100644 --- a/python/sglang/test/runners.py +++ b/python/sglang/test/runners.py @@ -12,10 +12,11 @@ # limitations under the License. # ============================================================================== +import json import multiprocessing as mp import os from dataclasses import dataclass -from typing import List, Optional, Tuple, Union +from typing import Any, List, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -89,7 +90,9 @@ def get_token_ids_logprobs(logits, token_ids): return logprobs -def _get_sentence_transformer_embedding_model(model_path, torch_dtype): +def _get_sentence_transformer_embedding_model( + model_path, torch_dtype, matryoshka_dim: Optional[int] = None +): from sentence_transformers import SentenceTransformer from sentence_transformers.util import is_sentence_transformer_model @@ -97,6 +100,7 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype): model = SentenceTransformer( model_path, model_kwargs={"torch_dtype": torch_dtype}, + truncate_dim=matryoshka_dim, ) else: # if no pre-trained sentence-transformers model from sentence_transformers import models @@ -106,7 +110,9 @@ def _get_sentence_transformer_embedding_model(model_path, torch_dtype): word_embedding_model.get_word_embedding_dimension(), pooling_mode="lasttoken", ) - model = SentenceTransformer(modules=[word_embedding_model, pooling_model]) + model = SentenceTransformer( + modules=[word_embedding_model, pooling_model], truncate_dim=matryoshka_dim + ) return model.cuda() @@ -135,6 +141,7 @@ class HFRunner: output_str_only: bool = False, trust_remote_code: bool = False, patch_model_do_sample_false: bool = False, + matryoshka_dim: Optional[int] = None, ): self.model_type = model_type self.output_str_only = output_str_only @@ -151,6 +158,7 @@ class HFRunner: self.out_queue, model_path, torch_dtype, + matryoshka_dim, ), ) self.model_proc.start() @@ -225,7 +233,14 @@ class HFRunner: embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1) return embeddings.contiguous() - def start_model_process(self, in_queue, out_queue, model_path, torch_dtype): + def start_model_process( + self, + in_queue, + out_queue, + model_path, + torch_dtype, + matryoshka_dim: Optional[int] = None, + ): # Apply model-specific patches monkey_patch_gemma2_sdpa() @@ -259,7 +274,7 @@ class HFRunner: self.processor = AutoProcessor.from_pretrained(model_path) else: self.model = _get_sentence_transformer_embedding_model( - model_path, torch_dtype + model_path, torch_dtype, matryoshka_dim=matryoshka_dim ) elif self.model_type == "reward" or self.model_type == "cross_encoder": from transformers import AutoModelForSequenceClassification @@ -519,6 +534,7 @@ class SRTRunner: lora_target_modules: Optional[List[str]] = None, enable_lora: Optional[bool] = None, max_loaded_loras: Optional[int] = None, + json_model_override_args: Optional[dict[str, Any]] = None, lora_eviction_policy: str = "lru", ): self.model_type = model_type @@ -566,6 +582,11 @@ class SRTRunner: lora_target_modules=lora_target_modules, enable_lora=enable_lora, max_loaded_loras=max_loaded_loras, + json_model_override_args=( + json.dumps(json_model_override_args) + if json_model_override_args + else "{}" + ), lora_eviction_policy=lora_eviction_policy, **spec_kwargs, ) @@ -594,6 +615,7 @@ class SRTRunner: logprob_start_len: int = 0, top_k: Optional[int] = None, token_ids_logprob: Optional[List[int]] = None, + dimensions: Optional[int] = None, ): if self.is_generation: return self.forward_generation_raw( @@ -607,7 +629,9 @@ class SRTRunner: ) else: if self.model_type == "embedding": - response = self.engine.encode(prompt=prompts, image_data=image_data) + response = self.engine.encode( + prompt=prompts, image_data=image_data, dimensions=dimensions + ) if isinstance(response, list): logits = [x["embedding"] for x in response] else: diff --git a/test/srt/models/test_embedding_models.py b/test/srt/models/test_embedding_models.py index c9dc86f1a..a93e762cf 100644 --- a/test/srt/models/test_embedding_models.py +++ b/test/srt/models/test_embedding_models.py @@ -15,6 +15,7 @@ import multiprocessing as mp import random import unittest +from typing import Optional import torch from transformers import AutoConfig, AutoTokenizer @@ -69,6 +70,7 @@ class TestEmbeddingModels(CustomTestCase): tp_size, torch_dtype, prefill_tolerance, + matryoshka_dim: Optional[int] = None, ) -> None: truncated_prompts = self._truncate_prompts(prompts, model_path) @@ -76,6 +78,7 @@ class TestEmbeddingModels(CustomTestCase): model_path, torch_dtype=torch_dtype, model_type="embedding", + matryoshka_dim=matryoshka_dim, ) as hf_runner: hf_outputs = hf_runner.forward(truncated_prompts) @@ -86,8 +89,13 @@ class TestEmbeddingModels(CustomTestCase): torch_dtype=torch_dtype, model_type="embedding", attention_backend=attention_backend, + json_model_override_args=( + {"matryoshka_dimensions": [matryoshka_dim]} if matryoshka_dim else None + ), ) as srt_runner: - srt_outputs = srt_runner.forward(truncated_prompts) + srt_outputs = srt_runner.forward( + truncated_prompts, dimensions=matryoshka_dim + ) for i in range(len(prompts)): hf_logits = torch.Tensor(hf_outputs.embed_logits[i]) @@ -113,6 +121,25 @@ class TestEmbeddingModels(CustomTestCase): DEFAULT_PROMPTS, model, tp_size, torch_dtype, prefill_tolerance ) + def test_matryoshka_embedding(self): + models_to_test = [ + model + for model in MODELS + if "Alibaba-NLP/gte-Qwen2-1.5B-instruct" == model[0] + ] + assert len(models_to_test) == 1 + + for model, tp_size, prefill_tolerance in models_to_test: + for torch_dtype in TORCH_DTYPES: + self.assert_close_prefill_logits( + DEFAULT_PROMPTS, + model, + tp_size, + torch_dtype, + prefill_tolerance, + matryoshka_dim=128, + ) + if __name__ == "__main__": unittest.main() diff --git a/test/srt/openai_server/basic/test_openai_embedding.py b/test/srt/openai_server/basic/test_openai_embedding.py index 60eb8f764..d198b1a7f 100644 --- a/test/srt/openai_server/basic/test_openai_embedding.py +++ b/test/srt/openai_server/basic/test_openai_embedding.py @@ -1,5 +1,8 @@ +import json +import os import unittest +import numpy as np import openai from sglang.srt.utils import kill_process_tree @@ -92,6 +95,105 @@ class TestOpenAIEmbedding(CustomTestCase): # check the status code self.assertEqual(cm.exception.status_code, 400) + def test_embedding_with_dimensions_parameter(self): + """Test that non-Matryoshka models reject dimensions parameter.""" + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + # Test that specifying dimensions fails for non-Matryoshka models + with self.assertRaises(openai.BadRequestError) as cm: + client.embeddings.create( + model=self.model, input="Hello world", dimensions=512 + ) + + self.assertEqual(cm.exception.status_code, 400) + + +class TestMatryoshkaEmbeddingModel(CustomTestCase): + """Test class for Model that supports Matryoshka embedding functionality, using OpenAI API.""" + + @classmethod + def setUpClass(cls): + cls.model = DEFAULT_SMALL_EMBEDDING_MODEL_NAME_FOR_TEST + cls.base_url = DEFAULT_URL_FOR_TEST + cls.api_key = "sk-123456" + cls.matryoshka_dims = [128, 256, 512, 768, 1024] + + # Configure embedding-specific args with Matryoshka support via json_model_override_args + matryoshka_config = { + "is_matryoshka": True, + "matryoshka_dimensions": cls.matryoshka_dims, + } + other_args = [ + "--is-embedding", + "--enable-metrics", + "--json-model-override-args", + json.dumps(matryoshka_config), + ] + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + api_key=cls.api_key, + other_args=other_args, + ) + cls.base_url += "/v1" + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process"): + kill_process_tree(cls.process.pid) + + def test_matryoshka_embedding_valid_dimensions(self): + """Test Matryoshka embedding with valid dimensions.""" + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + # Test with various valid dimensions + for dimensions in self.matryoshka_dims: + with self.subTest(dimensions=dimensions): + response = client.embeddings.create( + model=self.model, input="Hello world", dimensions=dimensions + ) + self.assertEqual(len(response.data), 1) + self.assertEqual(len(response.data[0].embedding), dimensions) + + def test_matryoshka_embedding_batch_same_dimensions(self): + """Test Matryoshka embedding with batch input and same dimensions.""" + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + response = client.embeddings.create( + model=self.model, + input=["Hello world", "Test text", "Another example"], + dimensions=256, + ) + + self.assertEqual(len(response.data), 3) + for embedding_data in response.data: + self.assertEqual(len(embedding_data.embedding), 256) + + def test_matryoshka_embedding_no_dimensions(self): + """Test embedding without specifying dimensions (should use full size).""" + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + response = client.embeddings.create(model=self.model, input="Hello world") + + self.assertEqual(len(response.data), 1) + + # Should return full embedding size when no dimensions specified + self.assertEqual(len(response.data[0].embedding), 1536) + + def test_matryoshka_embedding_invalid_dimensions(self): + """Test Matryoshka embedding with invalid dimensions.""" + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + for dimensions in [100, 0, -1, 10000]: + with self.assertRaises(openai.BadRequestError) as cm: + client.embeddings.create( + model=self.model, + input="Hello world", + dimensions=dimensions, + ) + self.assertEqual(cm.exception.status_code, 400) + if __name__ == "__main__": unittest.main()