weaviate
29 строк · 683.0 Байт
1// _ _
2// __ _____ __ ___ ___ __ _| |_ ___
3// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4// \ V V / __/ (_| |\ V /| | (_| | || __/
5// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6//
7// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8//
9// CONTACT: hello@weaviate.io
10//
11
12package clients
13
14import (
15"strings"
16
17"github.com/weaviate/tiktoken-go"
18)
19
20func GetTokensCount(model string, input string, tke *tiktoken.Tiktoken) int {
21tokensPerMessage := 3
22if strings.HasPrefix(model, "gpt-3.5-turbo") {
23tokensPerMessage = 4
24}
25
26tokensCount := tokensPerMessage
27tokensCount += len(tke.Encode(input, nil, nil))
28return tokensCount
29}
30