Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Token count fix (#2935)

Browse source 
Fix token count for OpenAIEmbeddings

Release Notes (Preview Only)

- update token count calculation for truncated OpenAIEmbeddings
- increased request timeout for OpenAI
This commit is contained in:
Kyle Caverly 2023-09-06 15:15:02 -04:00 • committed by Antonio Scandurra
parent 3b173b6090
commit f685265cad
1 changed files with 4 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

9
crates/semantic_index/src/embedding.rs
View file

@ -181,18 +181,17 @@ impl EmbeddingProvider for OpenAIEmbeddings {
fn truncate(&self, span: &str) -> (String, usize) { fn truncate(&self, span: &str) -> (String, usize) {
let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span); let mut tokens = OPENAI_BPE_TOKENIZER.encode_with_special_tokens(span);
let token_count = tokens.len(); let output = if tokens.len() > OPENAI_INPUT_LIMIT {
let output = if token_count > OPENAI_INPUT_LIMIT {
tokens.truncate(OPENAI_INPUT_LIMIT); tokens.truncate(OPENAI_INPUT_LIMIT);
OPENAI_BPE_TOKENIZER OPENAI_BPE_TOKENIZER
.decode(tokens) .decode(tokens.clone())
.ok() .ok()
.unwrap_or_else(|| span.to_string()) .unwrap_or_else(|| span.to_string())
} else { } else {
span.to_string() span.to_string()
}; };
(output, token_count) (output, tokens.len())
} }
async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> { async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
@ -204,7 +203,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
.ok_or_else(|| anyhow!("no api key"))?; .ok_or_else(|| anyhow!("no api key"))?;
let mut request_number = 0; let mut request_number = 0;
let mut request_timeout: u64 = 10; let mut request_timeout: u64 = 15;
let mut response: Response<AsyncBody>; let mut response: Response<AsyncBody>;
while request_number < MAX_RETRIES { while request_number < MAX_RETRIES {
response = self response = self