lookup-create.cpp 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. #include "arg.h"
  2. #include "common.h"
  3. #include "ngram-cache.h"
  4. #include "ggml.h"
  5. #include "llama.h"
  6. #include <cstdint>
  7. #include <fstream>
  8. #include <iostream>
  9. #include <string>
  10. #include <unordered_map>
  11. #include <vector>
  12. int main(int argc, char ** argv){
  13. gpt_params params;
  14. if (!gpt_params_parse(argc, argv, params, LLAMA_EXAMPLE_LOOKUP)) {
  15. return 1;
  16. }
  17. // init llama.cpp
  18. llama_backend_init();
  19. llama_numa_init(params.numa);
  20. // load the model
  21. llama_init_result llama_init = llama_init_from_gpt_params(params);
  22. llama_model * model = llama_init.model;
  23. llama_context * ctx = llama_init.context;
  24. GGML_ASSERT(model != nullptr);
  25. // tokenize the prompt
  26. std::vector<llama_token> inp;
  27. inp = ::llama_tokenize(ctx, params.prompt, true, true);
  28. fprintf(stderr, "%s: tokenization done\n", __func__);
  29. llama_ngram_cache ngram_cache;
  30. llama_ngram_cache_update(ngram_cache, LLAMA_NGRAM_STATIC, LLAMA_NGRAM_STATIC, inp, inp.size(), true);
  31. fprintf(stderr, "%s: hashing done, writing file to %s\n", __func__, params.lookup_cache_static.c_str());
  32. llama_ngram_cache_save(ngram_cache, params.lookup_cache_static);
  33. return 0;
  34. }