Repository: meta-llama/llama-stack
File: src/llama_stack/providers/utils/inference/embedding_mixin.py (Python)

Code before:
            data.append(
                OpenAIEmbeddingData(
                    embedding=embedding_value,
                    index=i,
                )
            )

        # Not returning actual token usage
        usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
        return OpenAIEmbeddingsResponse(
            data=data,
            model=params.model,
            usage=usage,
        )

    async def _load_sentence_transformer_model(self, model: str) -> "SentenceTransformer":
        global EMBEDDING_MODELS

        loaded_model = EMBEDDING_MODELS.get(model)
        if loaded_model is not None:
            return loaded_model

        log.info(f"Loading sentence transformer for {model}...")

        def _load_model():
            from sentence_transformers import SentenceTransformer

            platform_name = platform.system()
            if platform_name == DARWIN:
                # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
                # threads with the default parallel settings, so force a single-threaded CPU run.
                log.debug(f"Constraining torch threads on {platform_name} to a single worker")
                torch.set_num_threads(1)

            return SentenceTransformer(model, trust_remote_code=True)

        loaded_model = await asyncio.to_thread(_load_model)
        EMBEDDING_MODELS[model] = loaded_model
        return loaded_model


Review comment (suggestion):
let's load the model here rather than within the thread

```suggestion
                from sentence_transformers import SentenceTransformer
```

Code after:

            data.append(
                OpenAIEmbeddingData(
                    embedding=embedding_value,
                    index=i,
                )
            )

        # Not returning actual token usage
        usage = OpenAIEmbeddingUsage(prompt_tokens=-1, total_tokens=-1)
        return OpenAIEmbeddingsResponse(
            data=data,
            model=params.model,
            usage=usage,
        )

    async def _load_sentence_transformer_model(self, model: str) -> "SentenceTransformer":
        loaded_model = EMBEDDING_MODELS.get(model)
        if loaded_model is not None:
            return loaded_model

        async with EMBEDDING_MODELS_LOCK:
            loaded_model = EMBEDDING_MODELS.get(model)
            if loaded_model is not None:
                return loaded_model

            log.info(f"Loading sentence transformer for {model}...")

            def _load_model():
                from sentence_transformers import SentenceTransformer

                platform_name = platform.system()
                if platform_name == DARWIN:
                    # PyTorch's OpenMP kernels can segfault on macOS when spawned from background
                    # threads with the default parallel settings, so force a single-threaded CPU run.
                    log.debug(f"Constraining torch threads on {platform_name} to a single worker")
                    torch.set_num_threads(1)

                return SentenceTransformer(model, trust_remote_code=True)

            loaded_model = await asyncio.to_thread(_load_model)
            EMBEDDING_MODELS[model] = loaded_model
            return loaded_model