From 3f78e7c9c3f290176223c71c42232d1e465cca8e Mon Sep 17 00:00:00 2001 From: chottokun <29515187+chottokun@users.noreply.github.com> Date: Mon, 12 Oct 2026 00:34:00 +0000 Subject: [PATCH 1/2] perf: vectorize embedding formatting and encoding Vectorize Matryoshka dimensionality reduction, L2 re-normalization, and Base64 packing across embedding vector batches using NumPy/PyTorch. Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com> --- src/app/services/embedding.py | 104 ++++++++++++++++++++++------------ 1 file changed, 68 insertions(+), 36 deletions(-) diff --git a/src/app/services/embedding.py b/src/app/services/embedding.py index ad6d5e7..ac4a963 100644 --- a/src/app/services/embedding.py +++ b/src/app/services/embedding.py @@ -274,23 +274,58 @@ def _tokenize_and_truncate_embeddings( return processed_inputs, usage -def _format_embedding( - vector: List[float], dimensions: Optional[int], encoding_format: str -) -> Union[List[float], str]: - # Matryoshka dimensionality reduction - if dimensions is not None and dimensions > 0 and dimensions < len(vector): - vector = vector[:dimensions] - # L2 re-normalization - norm = math.sqrt(sum(x * x for x in vector)) - if norm > 0: - vector = [x / norm for x in vector] +def _format_embeddings_batch( + vectors: Any, dimensions: Optional[int], encoding_format: str +) -> List[Union[List[float], str]]: + import numpy as np + + if isinstance(vectors, (list, tuple)) and len(vectors) == 0: + return [] + + try: + import torch + + if isinstance(vectors, torch.Tensor): + arr = vectors.detach().cpu().numpy() + elif isinstance(vectors, np.ndarray): + arr = vectors + elif hasattr(vectors, "tolist") and callable(getattr(vectors, "tolist")): + arr = np.asarray(vectors.tolist()) + else: + arr = np.asarray(vectors) + except ImportError: + if isinstance(vectors, np.ndarray): + arr = vectors + elif hasattr(vectors, "tolist") and callable(getattr(vectors, "tolist")): + arr = np.asarray(vectors.tolist()) + else: + arr = np.asarray(vectors) + + if arr.size == 0: + return [] + + if arr.ndim == 1: + arr = arr[np.newaxis, :] + + # Matryoshka dimensionality reduction & L2 re-normalization + if dimensions is not None and 0 < dimensions < arr.shape[1]: + arr = arr[:, :dimensions] + arr_f64 = arr.astype(np.float64) + norms = np.linalg.norm(arr_f64, axis=1, keepdims=True) + norms[norms == 0] = 1.0 + arr = arr_f64 / norms if encoding_format == "base64": - # Pack list of floats as float32 little-endian binary ( Union[List[float], str]: + return _format_embeddings_batch([vector], dimensions, encoding_format)[0] class EmbeddingService(BaseEmbeddingService): @@ -344,14 +379,15 @@ async def create_embeddings(self, request: EmbeddingRequest) -> EmbeddingRespons {"input": processed_inputs, "model": request.model}, ) # Apply dimensions and encoding_format post-processing to TEI response - processed_data = [] - for item in data.get("data", []): - raw_emb = item["embedding"] - idx = item["index"] - formatted_emb = _format_embedding( - raw_emb, request.dimensions, request.encoding_format - ) - processed_data.append(EmbeddingData(embedding=formatted_emb, index=idx)) + tei_items = data.get("data", []) + raw_embs = [item["embedding"] for item in tei_items] + formatted_embs = _format_embeddings_batch( + raw_embs, request.dimensions, request.encoding_format + ) + processed_data = [ + EmbeddingData(embedding=formatted_emb, index=item["index"]) + for formatted_emb, item in zip(formatted_embs, tei_items) + ] data["data"] = [d.model_dump() for d in processed_data] return EmbeddingResponse(**data) @@ -440,14 +476,12 @@ def run_mm_inference(): return model.encode_multimodal(processed_items, **kwargs) embeddings = await anyio.to_thread.run_sync(run_mm_inference) + formatted_embs = _format_embeddings_batch( + embeddings, request.dimensions, request.encoding_format + ) response_data = [ - EmbeddingData( - embedding=_format_embedding( - emb, request.dimensions, request.encoding_format - ), - index=i, - ) - for i, emb in enumerate(embeddings) + EmbeddingData(embedding=emb, index=i) + for i, emb in enumerate(formatted_embs) ] usage = Usage(prompt_tokens=0, total_tokens=0) return EmbeddingResponse( @@ -474,14 +508,12 @@ def _run_inference(): vectors = await anyio.to_thread.run_sync(_run_inference) + formatted_embs = _format_embeddings_batch( + vectors, request.dimensions, request.encoding_format + ) response_data = [ - EmbeddingData( - embedding=_format_embedding( - vector, request.dimensions, request.encoding_format - ), - index=i, - ) - for i, vector in enumerate(vectors.tolist()) + EmbeddingData(embedding=emb, index=i) + for i, emb in enumerate(formatted_embs) ] return EmbeddingResponse(data=response_data, model=request.model, usage=usage) From d769e7d5331ee6169c4ea1500537f31a65e1afe8 Mon Sep 17 00:00:00 2001 From: chottokun Date: Mon, 12 Oct 2026 11:34:10 +0900 Subject: [PATCH 2/2] fix(lint): remove unused imports math and struct --- src/app/services/embedding.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/app/services/embedding.py b/src/app/services/embedding.py index 68e6b0c..6e4b84e 100644 --- a/src/app/services/embedding.py +++ b/src/app/services/embedding.py @@ -1,7 +1,5 @@ import asyncio import base64 -import math -import struct from typing import Any, List, Tuple, Optional, Union import anyio import httpx