forked from phoenix-oss/llama-stack-mirror
81ce39a607
We need to support: - asymmetric embedding models (#934) - truncation policies (#933) - varying dimensional output (#932) ## Test Plan ```bash $ cd llama_stack/providers/tests/inference $ pytest -s -v -k fireworks test_embeddings.py \ --inference-model nomic-ai/nomic-embed-text-v1.5 --env EMBEDDING_DIMENSION=784 $ pytest -s -v -k together test_embeddings.py \ --inference-model togethercomputer/m2-bert-80M-8k-retrieval --env EMBEDDING_DIMENSION=784 $ pytest -s -v -k ollama test_embeddings.py \ --inference-model all-minilm:latest --env EMBEDDING_DIMENSION=784 ```
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
from llama_stack.apis.inference import (
|
|
EmbeddingsResponse,
|
|
EmbeddingTaskType,
|
|
InterleavedContentItem,
|
|
ModelStore,
|
|
TextTruncation,
|
|
)
|
|
|
|
EMBEDDING_MODELS = {}
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class SentenceTransformerEmbeddingMixin:
|
|
model_store: ModelStore
|
|
|
|
async def embeddings(
|
|
self,
|
|
model_id: str,
|
|
contents: List[str] | List[InterleavedContentItem],
|
|
text_truncation: Optional[TextTruncation] = TextTruncation.none,
|
|
output_dimension: Optional[int] = None,
|
|
task_type: Optional[EmbeddingTaskType] = None,
|
|
) -> EmbeddingsResponse:
|
|
model = await self.model_store.get_model(model_id)
|
|
embedding_model = self._load_sentence_transformer_model(model.provider_resource_id)
|
|
embeddings = embedding_model.encode(contents)
|
|
return EmbeddingsResponse(embeddings=embeddings)
|
|
|
|
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}...")
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
loaded_model = SentenceTransformer(model)
|
|
EMBEDDING_MODELS[model] = loaded_model
|
|
return loaded_model
|