Add batching to embedding model

[wirthual]

PR to add batching to embedding model.

For models with tiktoken tokenizer, encode input and batch after this criteria:

• not a single input larger than 4096 token (doc says 8192 but it fails, 8191 works)
• the list of inputs not larger than 2048
• overall tokens not larger than 300_000

If a single input is longer than 4096 tokens, the input will be split into chunks and the final embedding is the average over the chunk embeddings (could be also weighted based on the length of the chunks)

Tokenized input is send to model.

Models which do not use tiktoken tokenizer continue to operate on strings.

Docs: https://platform.openai.com/docs/api-reference/embeddings/create#embeddings_create-input

Fixes #52.

[Guido van Rossum (gvanrossum)]

Ah, now I understand what you meant by averaging embeddings for long inputs. It'll take me a while to review this, it's the most substantial contribution ever! :-)

[Guido van Rossum (gvanrossum)]

Also could you work on the merge conflicts? Unfortunately I just did some work on the embedding code myself...

[Guido van Rossum (gvanrossum)]

Thanks! This is a good start.

I haven't reviewed everything yet, but I believe I've found a pretty big flaw (lists of ints) in the chunking code.

I also found a bunch of nits. And I think the code clarity would improve if everything was typed (something I try to enforce in my own code too).

[Guido van Rossum (gvanrossum)]

Looks unused?

[wirthual]

Removed.

[Guido van Rossum (gvanrossum)]

I find it hard to think about the code because entity_to_embed could be either a str or a list[int] (which encode() returns).

Moreover I think that in the case that it's a list[int], the rest of the code is wrong -- we end up sending lists of integers to the embedding endpoint, which AFAIK only expects lists of strings.

[wirthual]

The openAI endpoint also allows sending the tokenized inputs to get the embeddings. So this way it just skips the decoding part. However I understand this is a bit confusing and it might be clearer if the token chunks get decoded again and sent then.

[Guido van Rossum (gvanrossum)]

Oh, if sending tokens is supported, and we are getting tokens from tiktoken anyways, by all means send the tokens! I just didn't realize the API did that. (I looked in the docs but apparently in the wrong place -- do you have a deep link?)

[wirthual]

Yes, I got it from here

Input text to embed, encoded as a string or array of tokens.

[Guido van Rossum (gvanrossum)]

wirthual, I have another thought. Two actually.

First, the only place where we ever request embeddings of long inputs (controlled by the user) is when adding an index entry for a very long message to MessageTextIndex. However! We don't actually look up anything by embedding in that table. The embeddings are just used for a non-existent fallback if regular querying doesn't work, OR when explicitly using classic RAG to compare its effectiveness to structured RAG. But we don't have such a fallback and we don't do that comparison. (They were part of the original TypeScript code base and were somewhat mindlessly translated.)

So I don't think it's worth to complicate the code at all to deal with over-long single embedding input strings, except in order to make sure the entire batch isn't rejected by the embedding service. For this purpose simply truncating over-long input strings would seem good enough, and you won't need chunked_groups and any logic around it at all.

OTOH (this is the second thought) if you really, really, really want this embedding API completely rock solid for other use cases, can you research a bit what a typical classic RAG implementation does for over-long input strings?

[wirthual]

Hi,

Based on your description it makes sense to just make sure we do not run into the limits of the embedding service by truncating the inputs and limit the batch size. I can simplify this PR to account for that.

Generally splitting on sentence/section level rather than fixed size seems already a better approach. For example the implementation of RecursiveCharacterTextSplitter does this, where it recursively tries to split the text at certain characters, e.g. "\n\n", "\n", " ", "" to prevent text being cut off in the middle of an idea.

There are also more sophisticated methods which take the document type and structure or the semantics of the text into account e.g. with the help of an LLM.

[Guido van Rossum (gvanrossum)]

wirthual

Based on your description it makes sense to just make sure we do not run into the limits of the embedding service by truncating the inputs and limit the batch size. I can simplify this PR to account for that.

Sounds good!

Generally splitting on sentence/section level rather than fixed size seems already a better approach. For example the implementation of RecursiveCharacterTextSplitter does this, where it recursively tries to split the text at certain characters, e.g. "\n\n", "\n", " ", "" to prevent text being cut off in the middle of an idea.

Yeah, this should be left up to the process parsing the original input text into message chunks.

There are also more sophisticated methods which take the document type and structure or the semantics of the text into account e.g. with the help of an LLM.

Also thank you for reminding me that a lot of this stuff has already been done or at least been thought about. (Have you heard of something that can split email messages? That might help with #45...)

[Guido van Rossum (gvanrossum)]

Thanks -- great progress. I have a few nits only.

[Guido van Rossum (gvanrossum)]

LGTM.