Spaces:
Sleeping
Sleeping
| """ | |
| Transformation logic from OpenAI /v1/embeddings format to Azure AI Cohere's /v1/embed. | |
| Why separate file? Make it easy to see how transformation works | |
| Convers | |
| - Cohere request format | |
| Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html | |
| """ | |
| from typing import List, Optional, Tuple | |
| from litellm.types.llms.azure_ai import ImageEmbeddingInput, ImageEmbeddingRequest | |
| from litellm.types.llms.openai import EmbeddingCreateParams | |
| from litellm.types.utils import EmbeddingResponse, Usage | |
| from litellm.utils import is_base64_encoded | |
| class AzureAICohereConfig: | |
| def __init__(self) -> None: | |
| pass | |
| def _map_azure_model_group(self, model: str) -> str: | |
| if model == "offer-cohere-embed-multili-paygo": | |
| return "Cohere-embed-v3-multilingual" | |
| elif model == "offer-cohere-embed-english-paygo": | |
| return "Cohere-embed-v3-english" | |
| return model | |
| def _transform_request_image_embeddings( | |
| self, input: List[str], optional_params: dict | |
| ) -> ImageEmbeddingRequest: | |
| """ | |
| Assume all str in list is base64 encoded string | |
| """ | |
| image_input: List[ImageEmbeddingInput] = [] | |
| for i in input: | |
| embedding_input = ImageEmbeddingInput(image=i) | |
| image_input.append(embedding_input) | |
| return ImageEmbeddingRequest(input=image_input, **optional_params) | |
| def _transform_request( | |
| self, input: List[str], optional_params: dict, model: str | |
| ) -> Tuple[ImageEmbeddingRequest, EmbeddingCreateParams, List[int]]: | |
| """ | |
| Return the list of input to `/image/embeddings`, `/v1/embeddings`, list of image_embedding_idx for recombination | |
| """ | |
| image_embeddings: List[str] = [] | |
| image_embedding_idx: List[int] = [] | |
| for idx, i in enumerate(input): | |
| """ | |
| - is base64 -> route to image embeddings | |
| - is ImageEmbeddingInput -> route to image embeddings | |
| - else -> route to `/v1/embeddings` | |
| """ | |
| if is_base64_encoded(i): | |
| image_embeddings.append(i) | |
| image_embedding_idx.append(idx) | |
| ## REMOVE IMAGE EMBEDDINGS FROM input list | |
| filtered_input = [ | |
| item for idx, item in enumerate(input) if idx not in image_embedding_idx | |
| ] | |
| v1_embeddings_request = EmbeddingCreateParams( | |
| input=filtered_input, model=model, **optional_params | |
| ) | |
| image_embeddings_request = self._transform_request_image_embeddings( | |
| input=image_embeddings, optional_params=optional_params | |
| ) | |
| return image_embeddings_request, v1_embeddings_request, image_embedding_idx | |
| def _transform_response(self, response: EmbeddingResponse) -> EmbeddingResponse: | |
| additional_headers: Optional[dict] = response._hidden_params.get( | |
| "additional_headers" | |
| ) | |
| if additional_headers: | |
| # CALCULATE USAGE | |
| input_tokens: Optional[str] = additional_headers.get( | |
| "llm_provider-num_tokens" | |
| ) | |
| if input_tokens: | |
| if response.usage: | |
| response.usage.prompt_tokens = int(input_tokens) | |
| else: | |
| response.usage = Usage(prompt_tokens=int(input_tokens)) | |
| # SET MODEL | |
| base_model: Optional[str] = additional_headers.get( | |
| "llm_provider-azureml-model-group" | |
| ) | |
| if base_model: | |
| response.model = self._map_azure_model_group(base_model) | |
| return response | |