tokenize.cpp 1.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344
  1. #include "common.h"
  2. #include "llama.h"
  3. #include <cmath>
  4. #include <cstdio>
  5. #include <string>
  6. #include <vector>
  7. int main(int argc, char ** argv) {
  8. if (argc < 3 || argv[1][0] == '-') {
  9. printf("usage: %s MODEL_PATH PROMPT [--ids]\n" , argv[0]);
  10. return 1;
  11. }
  12. const char * model_path = argv[1];
  13. const char * prompt = argv[2];
  14. const bool printing_ids = argc > 3 && std::string(argv[3]) == "--ids";
  15. llama_backend_init();
  16. llama_model_params model_params = llama_model_default_params();
  17. model_params.vocab_only = true;
  18. llama_model * model = llama_load_model_from_file(model_path, model_params);
  19. llama_context_params ctx_params = llama_context_default_params();
  20. llama_context * ctx = llama_new_context_with_model(model, ctx_params);
  21. const bool add_bos = llama_should_add_bos_token(model);
  22. std::vector<llama_token> tokens;
  23. tokens = ::llama_tokenize(model, prompt, add_bos, true);
  24. for (int i = 0; i < (int) tokens.size(); i++) {
  25. if (printing_ids) {
  26. printf("%d\n", tokens[i]);
  27. } else {
  28. printf("%6d -> '%s'\n", tokens[i], llama_token_to_piece(ctx, tokens[i]).c_str());
  29. }
  30. }
  31. return 0;
  32. }