Skip to content
This repository was archived by the owner on Sep 30, 2024. It is now read-only.

Add Support for Counting Tokens for Azure Code and Update in Redis#63100

Merged
arafatkatze merged 2 commits into
mainfrom
add-azure-tokenization
Jun 28, 2024
Merged

Add Support for Counting Tokens for Azure Code and Update in Redis#63100
arafatkatze merged 2 commits into
mainfrom
add-azure-tokenization

Conversation

@arafatkatze

@arafatkatze arafatkatze commented Jun 5, 2024

Copy link
Copy Markdown
Contributor

Description:

This PR introduces support for counting tokens within the Azure code and updating these counts in Redis. The token counting logic is embedded directly in the Azure code rather than using a standardized point for all token counting logic.

Reasoning:

•	Azure does not currently support obtaining token usage from their streaming endpoint, unlike OpenAI.
•	To enable immediate functionality, the token counting logic is placed within the Azure code itself.
•	The implementation supports GPT-4o.

Future Considerations:

•	When Azure eventually adds support for token usage from the streaming endpoint, we will migrate to using Azure’s built-in capabilities.
•	This will ensure full utilization of Azure OpenAI features as they achieve parity with OpenAI.

Changes:

•	Added token counting logic to the Azure code.
•	Updated Redis with the token counts.

Testing:

•	Verified the implementation works with GPT-4o.

Conclusion:

This is a temporary solution to enable token counting in Azure. We will adapt our approach as Azure enhances its feature set to include token usage from their streaming endpoint.

Test plan

Tested locally with debugger

Changelog

@cla-bot cla-bot Bot added the cla-signed label Jun 5, 2024
@arafatkatze arafatkatze changed the title Add azure tokenization and token counting support Add Support for Counting Tokens for Azure Code and Update in Redis Jun 5, 2024
@arafatkatze arafatkatze requested review from emidoots June 5, 2024 13:30
Comment thread schema/site.schema.json
"gpt-4"
]
},
"azureCompletionModel": {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I hate adding extra parameters to the siteconfig and requestparams. However, I had to add these parameters because Azure does not support naming the model directly. Instead, you provide the deployment name. The deployment name can be different from the model name. Because of this discrepancy, I had to introduce extra config parameters to specify the chat model name and the completion model name.

This necessity arises because Azure doesn’t allow retrieval of the model name straightforwardly. The deployment name, which is used in requests, can map to any underlying model. For accurate token counting, the correct model name is required. Although there are ways to retrieve the model name from the deployment, they involve complex logic. Therefore, this approach, despite being somewhat inelegant, is the simplest solution given the constraints.

Comment thread go.mod
github.com/pkoukk/tiktoken-go v0.1.6
github.com/pkoukk/tiktoken-go-loader v0.0.1
github.com/pkoukk/tiktoken-go v0.1.7
github.com/pkoukk/tiktoken-go-loader v0.0.2-0.20240522064338-c17e8bc0f699

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Had to update to support gpt4o

@emidoots

emidoots commented Jun 6, 2024

Copy link
Copy Markdown
Member

Will review this tomorrow, apologies for the delay

Comment thread internal/completions/client/azureopenai/openai.go Outdated
@arafatkatze

arafatkatze commented Jun 6, 2024

Copy link
Copy Markdown
Contributor Author

Added the validation check also in this commit
image

@arafatkatze arafatkatze force-pushed the add-azure-tokenization branch from 561199a to 52b92b9 Compare June 6, 2024 20:04
Comment thread internal/completions/client/azureopenai/openai.go Outdated
Comment thread internal/conf/validate_custom.go Outdated
@emidoots

Copy link
Copy Markdown
Member

two minor comments, will be LGTM once those are addressed

@arafatkatze arafatkatze force-pushed the add-azure-tokenization branch from 52b92b9 to 5ba23c8 Compare June 27, 2024 12:45

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a comment above this line:

// Note: If we had an error calculating input/output tokens, that is unfortunate, the
// best thing we can do is record zero token usage which would be our hint to look at
// the logs for errors.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same comment here

@arafatkatze arafatkatze enabled auto-merge (squash) June 27, 2024 16:12
@arafatkatze arafatkatze force-pushed the add-azure-tokenization branch from dc2c248 to aff8378 Compare June 28, 2024 11:41
@arafatkatze arafatkatze force-pushed the add-azure-tokenization branch from aff8378 to 0dcdf85 Compare June 28, 2024 11:41
@arafatkatze arafatkatze merged commit 141d2e0 into main Jun 28, 2024
@arafatkatze arafatkatze deleted the add-azure-tokenization branch June 28, 2024 12:37
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants