File size: 3,621 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
"""
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
|