Improve time complexity of '_byte_pair_merge' and add test

[michaelgiba]

I saw the comment in the byte_pair_merge code about improving the time complexity and thought it sounded interesting. I'm interested to see how the overall throughput of tiktoken changes with this change. I'll work on benchmarking it when I get a chance

[hauntsaninja]

Thanks for looking into this! I'd played around with it too a long time back and found the constant factor was dominant in practice. Curious to see the benchmarking results here; we could also consider conditioning on piece length.

One quick note about the tests: I haven't gotten around to open sourcing most of tiktoken's tests. So apologies if that made your life harder.

[michaelgiba]

No problem on the tests, I figured that might be the case. Are there more benchmarking tools internally as well or is that something I could contribute organically ? If the piece lengths are typically small I would imagine this solution is slower given the extra upfront cost.

…

On Tue, Feb 7, 2023 at 7:45 PM Shantanu ***@***.***> wrote: Thanks for looking into this! I'd played around with it too a long time back and found the constant factor was dominant in practice. Curious to see the benchmarking results here; we could also consider conditioning on piece length. One quick note about the tests: I haven't gotten around to open sourcing most of tiktoken's tests. So apologies if that made your life harder. — Reply to this email directly, view it on GitHub <#29 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABTFWYOYFVZSVJ7557HPORTWWL3DPANCNFSM6AAAAAAUUVEDFQ> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[hauntsaninja]

There is some more internal benchmarking, but it looks a lot like https://github.com/openai/tiktoken/blob/main/scripts/benchmark.py

The main thing is just finding a documents: list[str] to feed to it. I used some internal-only datasets for this purpose, but maybe openwebtext or something would work. 100MB of data should be plenty for a meaningful benchmark (at low thread count).

Piece lengths are typically very small (about word length). That said, there are some degenerate cases. The best case scenario for heap merges is probably something like:

import base64, random

def base64_noise_documents(n_docs: int):
    rand = random.Random(217)
    documents = []
    for _ in range(n_docs):
        documents.append(base64.b64encode(rand.randbytes(rand.randint(100, 10_000))).decode())
    return documents

[michaelgiba]

Ok I had a chance to run some simple tests. I ran with 1,3,5,7,9,11 threads on both the original version and the heap version against documents created using your random b64 example above and 100 document samples from this random ubuntu dialog dataset I found https://github.com/rkadlec/ubuntu-ranking-dataset-creator

Surprisingly the performance looks decent. I'll try to run something more rigorous when I have free time.

[michaelgiba]

Haven't gotten around to testing this - but I'm closing it out in favor of #31 which looks great!

[hauntsaninja]

Thanks for experimenting with it though! :-)