Properly return usage information for BedrockClientV2 for the Reranker model?

[ldorigo]

It would be very useful to have a number of "used billable units" since now I need to implement some complicated (and cpu-intensive) logic to tokenize the queries, tokenize each document; and count how many billable requests will be used by a given request based on the amount and length of documents.

[ldorigo]

Hi @billytrend-cohere , if at all possible, could you indicate whether this is planned and perhaps give a rough ETA? We found that doing the tokenization ourselves is indeed quite resource-intensive at scale (and increased memory usage quite significantly); I need to decide whether it makes sense to try and optimize it on our end or wait for it to be supported.

[billytrend-cohere]

Hey @ldorigo Looking into this for you! you can use our platform to do free tokenisation if you need to be unblocked sooner. just set offline=False on the tokenise call

[rehan243]

Fwiw, we hit a similar issue with our RAG pipeline using the Cohere Python client v4.2.1. To get around it, we ended up implementing a custom tokenizer to estimate billable units, but it's been a pain to maintain. We can simplify this by directly returning usage info in BedrockClientV2 like this:

response = client.rerank(...)
usage = {'billable_units': response.meta['billed_units']}

This way, we don't have to worry about tokenizing queries and docs manually. We've seen significant CPU savings since implementing a similar fix for our embedding models.