GHghita_·Dec 4, 2024Submission
Show HN: I Used Llama-70B Logprobs for Better, Cheaper and Faster Chunking
The LlamaChunk algorithm is simple: We pick a special token that we know is not in the corpus, e.g. "段".
We pick the character "段" because
1. tiktoken always encodes it to exactly one token,
2. it is not in the corpus, and
3. it means "section" in Chinese.
Then, we ask Llama to repeat the User's message with the "段" token sprinkled throughout.
And just like that, it's perfect out of the box! It correctly used the "段" character to mark the end of every section.
BUT if you've ever worked with LLM's, you know that input tokens are processed almost instantly, and output tokens take an eternity to generate. A naïve method is to simply wait for the LLM to repeat the entire python code, inserting "段" throughout.
However, by inferencing llama locally, we have a vastly more efficient way of doing this! We can simply pass in the entire paragraph
Read the rest
, and check the logprobs to see the probability that Llama wanted to output a "段" token at that location!
Done! The high logprob values, clearly indicate the locations that we should chunk! And, this is only possible because we have direct low-level access to Llama 3.1 70B.
Of course, there is a caveat. Because there are no output tokens, Llama can no longer see its own line breaks. Thus, as the text gets longer, it loses the willpower to continue to want to output "段"
But, we can simply normalize by this decaying curve in order to to fix this. And now, we're ready to split any type of document, without having to resort to regex or manually created rules.
Processing 450,000 characters took about 15 minutes on an A100. However, ~80% of that time was saving and loading the KVCache (Which can be done instantly if written in C++, rather than Python). So, we can expect that it would take 3 minutes per 450,000 characters if done optimally. Or, 7 MTokens per hour.
In terms of accuracy, LlamaChunk has higher recall AND precision, than a naïve chunking method and than semantic chunking (Which uses embeddings to detect sentence split boundaries, and still requires a good regex-based sentence splitter).
---
One thing you might wonder: What if the ideal chunk split is not along a token boundary?
Well for one thing, this is rare, as tokenizers intentionally split along meaningful boundaries. However, if the best split really is after the "f" in "fun", then you calculate
```
lp = logprob(prefix="Fine Tuning is", next_token=" f")
if lp > -7:
lp *= logprob(prefix="Fine Tuning is f", next_token="段")
```
In other words, if the logprob of token " f" has a non-negligible value > -7, then we can multiply by the logprob of the token after that being "段". The first line of code has a prefix that matches the main inference, so it does not need to be recalculated. However, if the "if" statement passes, then we'll have to do an extra inference, costing us latency. However, this "if" statement almost never passes (In our measurement, it happens once every ~2000 tokens, so it amortizes well).
Check out the repo for examples and more technical details.