Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use String, Dict, and read_bytes to shorten and simplify #91

Merged
merged 5 commits into from
May 21, 2024

Conversation

mikowals
Copy link
Contributor

This is based off current nightly branch (mojo 2024.4.161). It is a demo of some clean ups that can happen now that Mojo and its stdlib have added a lot of functionality that was missing when this was originally released. There could probably also be another round to remove TensorSlice and just use List[TensorF32] for each layer of weights.

The main changes are:

  • replace PointerString and PointerStrings with String and List[String] respectively.
  • Use a Dict to lookup token indices. This allows removing Quicksort and binary search code.
  • small changes like using Tensor's own SIMD operations and argmax.
  • use read_bytes to handling of pointers without copying. That means FileBuf is no longer used.

I am not sure this is ready to merge mostly because the of handling of special bytes handling in wrap and the old print function. I tried to persevere the functionality but haven't tested extensively. Ideally we could get proper handling from String and if not fix it in stdlib.

Also, I think the stdlib is going to shift to List[UInt8] for all bytes representations, including in String. So this change could also wait until after has happened and is incorporated.

I didn't mess with Llamatune since this is going across Mojo versions but locally there was no change in tokens / sec. It is probably loading faster and more memory efficiently since this avoids the vocab sort and no longer reads entire tokenizer.bin.

llama2.mojo Outdated
right = mid - 1
var index = self.map_vocab_to_index.find(token)
if index:
return index.value()

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

        return index.or_else(-1)

@tairov tairov marked this pull request as ready for review May 21, 2024 10:33
@tairov tairov merged commit f41dde4 into tairov:master May 21, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants