aws-samples / bedrock-access-gateway

OpenAI-Compatible RESTful APIs for Amazon Bedrock
MIT No Attribution
124 stars 24 forks source link

feat: Added support for Cohere Embed and Titan Embeddings models #3

Closed JGalego closed 3 months ago

JGalego commented 3 months ago

This PR adds support for Bedrock-provided Embeddings models (text only) by implementing the OpenAI's embeddings route.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

Example:

import os
from openai import OpenAI

client = OpenAI(
    base_url=os.environ['OPENAI_BASE_URL'],
    api_key=os.environ['OPENAI_API_KEY']
)

##### Models

models = client.models.list()

embedding_models = list(
    map(
        lambda model: model.id,
        filter(
            lambda model: "embed" in model.id,
            models.data
        )
))

print(embedding_models)

##### Embeddings

embeddings = client.embeddings.create(
    input = ["Olá mundo!", "Hello World!"],
    model="cohere.embed-multilingual-v3"
)

for item in embeddings.data:
    print(item.embedding[:10])

Output:

['cohere.embed-multilingual-v3', 'cohere.embed-english-v3', 'amazon.titan-embed-text-v1', 'amazon.titan-embed-image-v1']
[0.0017204285, 0.04119873, 0.024536133, -0.01625061, 0.001581192, 0.010940552, -0.0211792, -0.050720215, -0.041534424, 0.0053710938]
[0.008758545, 0.050598145, 0.022613525, 0.0124435425, -0.024017334, -0.0057029724, -0.03451538, -0.043151855, -0.039520264, 0.034698486]