diff --git a/.gitignore b/.gitignore index 072945180..9fb5b80c3 100644 --- a/.gitignore +++ b/.gitignore @@ -77,6 +77,7 @@ models-mnt /batched-bench /export-lora /finetune +/retrieval /speculative /parallel /train-text-from-scratch diff --git a/Makefile b/Makefile index 4f260cc3d..130fde838 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ BUILD_TARGETS = \ main quantize quantize-stats perplexity imatrix embedding vdot q8dot train-text-from-scratch convert-llama2c-to-ggml \ simple batched batched-bench save-load-state server gguf gguf-split llama-bench libllava.a llava-cli baby-llama beam-search \ - speculative infill tokenize benchmark-matmult parallel finetune export-lora lookahead lookup passkey gritlm tests/test-c.o + retrieval speculative infill tokenize benchmark-matmult parallel finetune export-lora lookahead lookup passkey gritlm tests/test-c.o # Binaries only useful for tests TEST_TARGETS = \ @@ -804,6 +804,10 @@ export-lora: examples/export-lora/export-lora.cpp ggml.o common/common.h $(OBJS) $(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<) $(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS) +retrieval: examples/retrieval/retrieval.cpp ggml.o llama.o $(COMMON_DEPS) $(OBJS) + $(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<) + $(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS) + speculative: examples/speculative/speculative.cpp ggml.o llama.o $(COMMON_DEPS) grammar-parser.o $(OBJS) $(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<) $(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS) diff --git a/common/common.cpp b/common/common.cpp index fb80d4bf7..9dec08430 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -157,7 +157,7 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) { return result; } -static bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_params & params, int & i, bool & invalid_param) { +bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_params & params, int & i, bool & invalid_param) { llama_sampling_params& sparams = params.sparams; if (arg == "-s" || arg == "--seed") { diff --git a/common/common.h b/common/common.h index a223eceaa..99ee90bc3 100644 --- a/common/common.h +++ b/common/common.h @@ -171,6 +171,8 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params); void gpt_print_usage(int argc, char ** argv, const gpt_params & params); +bool gpt_params_find_arg(int argc, char ** argv, const std::string & arg, gpt_params & params, int & i, bool & invalid_param); + std::string get_system_info(const gpt_params & params); std::string gpt_random_prompt(std::mt19937 & rng); diff --git a/common/log.h b/common/log.h index eb111e784..48d21e43c 100644 --- a/common/log.h +++ b/common/log.h @@ -566,6 +566,7 @@ inline void log_print_usage() printf(" --log-new Create a separate new log file on start. " "Each log file will have unique name: \"..log\"\n"); printf(" --log-append Don't truncate the old log file.\n"); + printf("\n"); } #define log_dump_cmdline(argc, argv) log_dump_cmdline_impl(argc, argv) diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index b59cc65bf..76496bf06 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -34,6 +34,7 @@ else() add_subdirectory(perplexity) add_subdirectory(quantize) add_subdirectory(quantize-stats) + add_subdirectory(retrieval) add_subdirectory(save-load-state) add_subdirectory(simple) add_subdirectory(passkey) diff --git a/examples/retrieval/CMakeLists.txt b/examples/retrieval/CMakeLists.txt new file mode 100644 index 000000000..eaabae08d --- /dev/null +++ b/examples/retrieval/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET retrieval) +add_executable(${TARGET} retrieval.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_11) diff --git a/examples/retrieval/README.md b/examples/retrieval/README.md new file mode 100644 index 000000000..2b2595c46 --- /dev/null +++ b/examples/retrieval/README.md @@ -0,0 +1,69 @@ +# llama.cpp/examples/retrieval + +Demonstration of simple retrieval technique based on cosine similarity + +More info: +https://github.com/ggerganov/llama.cpp/pull/6193 + +### How to use + +`retieval.cpp` has parameters of its own: +- `--context-file`: file to be embedded - state this option multiple times to embed multiple files +- `--chunk-size`: minimum size of each text chunk to be embedded +- `--chunk-separator`: STRING to divide chunks by. newline by default + +`retrieval` example can be tested as follows: + +```bash +make -j && ./retrieval --model ./models/bge-base-en-v1.5-f16.gguf --top-k 3 --context-file README.md --context-file License --chunk-size 100 --chunk-separator . +``` + +This chunks and embeds all given files and starts a loop requesting query inputs: + +``` +Enter query: +``` + +On each query input, top k chunks are shown along with file name, chunk position within file and original text: + +``` +Enter query: describe the mit license +batch_decode: n_tokens = 6, n_seq = 1 +Top 3 similar chunks: +filename: README.md +filepos: 119 +similarity: 0.762334 +textdata: +png) + +[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT) + +[Roadmap](https://github. +-------------------- +filename: License +filepos: 0 +similarity: 0.725146 +textdata: +MIT License + +Copyright (c) 2023 Georgi Gerganov + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +-------------------- +filename: README.md +filepos: 9178 +similarity: 0.621722 +textdata: +com/cztomsik/ava) (MIT) +- [ptsochantaris/emeltal](https://github.com/ptsochantaris/emeltal) +- [pythops/tenere](https://github. +-------------------- +``` diff --git a/examples/retrieval/retrieval.cpp b/examples/retrieval/retrieval.cpp new file mode 100644 index 000000000..5ba71e76a --- /dev/null +++ b/examples/retrieval/retrieval.cpp @@ -0,0 +1,350 @@ +#include "common.h" +#include "llama.h" + +#include +#include + +struct retrieval_params { + std::vector context_files; // context files to embed + int32_t chunk_size = 64; // chunk size for context embedding + std::string chunk_separator = "\n"; // chunk separator for context embedding +}; + +static void retrieval_params_print_usage(int argc, char ** argv, gpt_params & gpt_params, retrieval_params & params) { + gpt_print_usage(argc, argv, gpt_params); + printf("retrieval options:\n"); + printf(" --context-file FNAME file containing context to embed.\n"); + printf(" specify multiple files by providing --context-file option multiple times.\n"); + printf(" --chunk-size N minimum length of embedded text chunk (default:%d)\n", params.chunk_size); + printf(" --chunk-separator STRING\n"); + printf(" string to separate chunks (default: \"\\n\")\n"); + printf("\n"); +} + +static void retrieval_params_parse(int argc, char ** argv, gpt_params & gpt_params, retrieval_params & retrieval_params) { + int i = 1; + std::string arg; + while (i < argc) { + arg = argv[i]; + bool invalid_gpt_param = false; + if(gpt_params_find_arg(argc, argv, argv[i], gpt_params, i, invalid_gpt_param)) { + if (invalid_gpt_param) { + fprintf(stderr, "error: invalid argument: %s\n", arg.c_str()); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + // option was parsed by gpt_params_find_arg + } else if (arg == "--context-file") { + if (++i >= argc) { + fprintf(stderr, "error: missing argument for --context-file\n"); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + std::ifstream file(argv[i]); + if (!file) { + fprintf(stderr, "error: failed to open file '%s'\n", argv[i]); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + // store the external file name in params + retrieval_params.context_files.push_back(argv[i]); + } else if (arg == "--chunk-size") { + if (++i >= argc) { + fprintf(stderr, "error: missing argument for --chunk-size\n"); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + retrieval_params.chunk_size = std::stoi(argv[i]); + } else if (arg == "--chunk-separator") { + if (++i >= argc) { + fprintf(stderr, "error: missing argument for --chunk-separator\n"); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + retrieval_params.chunk_separator = argv[i]; + } else { + // unknown argument + fprintf(stderr, "error: unknown argument: %s\n", arg.c_str()); + retrieval_params_print_usage(argc, argv, gpt_params, retrieval_params); + exit(1); + } + i++; + } +} + +struct chunk { + // filename + std::string filename; + // original file position + size_t filepos; + // original text data + std::string textdata = ""; + // tokenized text data + std::vector tokens; + // embedding + std::vector embedding; +}; + +// chunk file data to chunks of size >= chunk_size +// chunk_separator is the separator between chunks +static std::vector chunk_file(const std::string & filename, int chunk_size, const std::string & chunk_separator) { + std::vector chunks; + std::ifstream f(filename.c_str()); + + if (!f.is_open()) { + fprintf(stderr, "Error: could not open file %s\n", filename.c_str()); + return chunks; + } + + chunk current_chunk; + char buffer[1024]; + int64_t filepos = 0; + std::string current = ""; + while (f.read(buffer, 1024)) { + current += std::string(buffer, f.gcount()); + size_t pos; + while ((pos = current.find(chunk_separator)) != std::string::npos) { + current_chunk.textdata += current.substr(0, pos + chunk_separator.size()); + if ((int) current_chunk.textdata.size() > chunk_size) { + // save chunk + current_chunk.filepos = filepos; + current_chunk.filename = filename; + chunks.push_back(current_chunk); + // update filepos + filepos += (int) current_chunk.textdata.size(); + // reset current_chunk + current_chunk = chunk(); + } + current = current.substr(pos + chunk_separator.size()); + } + + } + // add leftover data to last chunk + if (current_chunk.textdata.size() > 0) { + if (chunks.empty()) { + current_chunk.filepos = filepos; + current_chunk.filename = filename; + chunks.push_back(current_chunk); + } else { + chunks.back().textdata += current_chunk.textdata; + } + } + f.close(); + return chunks; +} + +static void batch_add_seq(llama_batch & batch, const std::vector & tokens, int seq_id) { + for (size_t i = 0; i < tokens.size(); i++) { + llama_batch_add(batch, tokens[i], i, { seq_id }, i == tokens.size() - 1); + } +} + +static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) { + // clear previous kv_cache values (irrelevant for embeddings) + llama_kv_cache_clear(ctx); + + // run model + fprintf(stderr, "%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq); + if (llama_decode(ctx, batch) < 0) { + fprintf(stderr, "%s : failed to decode\n", __func__); + } + + for (int i = 0; i < batch.n_tokens; i++) { + if (!batch.logits[i]) { + continue; + } + + // try to get sequence embeddings - supported only when pooling_type is not NONE + const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]); + if (embd == NULL) { + embd = llama_get_embeddings_ith(ctx, i); + if (embd == NULL) { + fprintf(stderr, "%s: failed to get embeddings for token %d\n", __func__, i); + continue; + } + } + + float * out = output + batch.seq_id[i][0] * n_embd; + llama_embd_normalize(embd, out, n_embd); + } +} + +int main(int argc, char ** argv) { + gpt_params params; + retrieval_params retrieval_params; + + retrieval_params_parse(argc, argv, params, retrieval_params); + + // For BERT models, batch size must be equal to ubatch size + params.n_ubatch = params.n_batch; + + if (retrieval_params.chunk_size <= 0) { + fprintf(stderr, "chunk_size must be positive\n"); + return 1; + } + if (retrieval_params.context_files.empty()) { + fprintf(stderr, "context_files must be specified\n"); + return 1; + } + params.embedding = true; + + print_build_info(); + + printf("processing files:\n"); + for (auto & context_file : retrieval_params.context_files) { + printf("%s\n", context_file.c_str()); + } + + std::vector chunks; + for (auto & context_file : retrieval_params.context_files) { + std::vector file_chunk = chunk_file(context_file, retrieval_params.chunk_size, retrieval_params.chunk_separator); + chunks.insert(chunks.end(), file_chunk.begin(), file_chunk.end()); + } + printf("Number of chunks: %ld\n", chunks.size()); + + llama_backend_init(); + llama_numa_init(params.numa); + + llama_model * model; + llama_context * ctx; + + // load the model + std::tie(model, ctx) = llama_init_from_gpt_params(params); + if (model == NULL) { + fprintf(stderr, "%s: error: unable to load model\n", __func__); + return 1; + } + + const int n_ctx_train = llama_n_ctx_train(model); + const int n_ctx = llama_n_ctx(ctx); + + if (n_ctx > n_ctx_train) { + fprintf(stderr, "%s: warning: model was trained on only %d context tokens (%d specified)\n", + __func__, n_ctx_train, n_ctx); + } + + // print system information + { + fprintf(stderr, "\n"); + fprintf(stderr, "%s\n", get_system_info(params).c_str()); + } + + // max batch size + const uint64_t n_batch = params.n_batch; + GGML_ASSERT(params.n_batch >= params.n_ctx); + + // tokenize the prompts and trim + for (auto & chunk : chunks) { + auto inp = ::llama_tokenize(ctx, chunk.textdata, true, false); + if (inp.size() > n_batch) { + fprintf(stderr, "%s: error: chunk size (%lld) exceeds batch size (%lld), increase batch size and re-run\n", + __func__, (long long int) inp.size(), (long long int) n_batch); + return 1; + } + // add eos if not present + if (inp.empty() || inp.back() != llama_token_eos(model)) { + inp.push_back(llama_token_eos(model)); + } + chunk.tokens = inp; + } + + // tokenization stats + if (params.verbose_prompt) { + for (int i = 0; i < (int) chunks.size(); i++) { + fprintf(stderr, "%s: prompt %d: '%s'\n", __func__, i, chunks[i].textdata.c_str()); + fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, chunks[i].tokens.size()); + for (int j = 0; j < (int) chunks[i].tokens.size(); j++) { + fprintf(stderr, "%6d -> '%s'\n", chunks[i].tokens[j], llama_token_to_piece(ctx, chunks[i].tokens[j]).c_str()); + } + fprintf(stderr, "\n\n"); + } + } + + // initialize batch + const int n_chunks = chunks.size(); + struct llama_batch batch = llama_batch_init(n_batch, 0, 1); + + // allocate output + const int n_embd = llama_n_embd(model); + std::vector embeddings(n_chunks * n_embd, 0); + float * emb = embeddings.data(); + + // break into batches + int p = 0; // number of prompts processed already + int s = 0; // number of prompts in current batch + for (int k = 0; k < n_chunks; k++) { + // clamp to n_batch tokens + auto & inp = chunks[k].tokens; + + const uint64_t n_toks = inp.size(); + + // encode if at capacity + if (batch.n_tokens + n_toks > n_batch) { + float * out = emb + p * n_embd; + batch_decode(ctx, batch, out, s, n_embd); + llama_batch_clear(batch); + p += s; + s = 0; + } + + // add to batch + batch_add_seq(batch, inp, s); + s += 1; + } + + // final batch + float * out = emb + p * n_embd; + batch_decode(ctx, batch, out, s, n_embd); + + // save embeddings to chunks + for (int i = 0; i < n_chunks; i++) { + chunks[i].embedding = std::vector(emb + i * n_embd, emb + (i + 1) * n_embd); + // clear tokens as they are no longer needed + chunks[i].tokens.clear(); + } + + // start loop, receive query and return top k similar chunks based on cosine similarity + std::string query; + while (true) { + printf("Enter query: "); + std::getline(std::cin, query); + std::vector query_tokens = llama_tokenize(ctx, query, true); + + struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1); + batch_add_seq(query_batch, query_tokens, 0); + + std::vector query_emb(n_embd, 0); + batch_decode(ctx, query_batch, query_emb.data(), 1, n_embd); + + llama_batch_clear(query_batch); + + // compute cosine similarities + { + std::vector> similarities; + for (int i = 0; i < n_chunks; i++) { + float sim = llama_embd_similarity_cos(chunks[i].embedding.data(), query_emb.data(), n_embd); + similarities.push_back(std::make_pair(i, sim)); + } + + // sort similarities + std::sort(similarities.begin(), similarities.end(), [](const std::pair & a, const std::pair & b) { + return a.second > b.second; + }); + + printf("Top %d similar chunks:\n", params.sparams.top_k); + for (int i = 0; i < std::min(params.sparams.top_k, (int) chunks.size()); i++) { + printf("filename: %s\n", chunks[similarities[i].first].filename.c_str()); + printf("filepos: %lld\n", (long long int) chunks[similarities[i].first].filepos); + printf("similarity: %f\n", similarities[i].second); + printf("textdata:\n%s\n", chunks[similarities[i].first].textdata.c_str()); + printf("--------------------\n"); + } + } + } + + // clean up + llama_print_timings(ctx); + llama_free(ctx); + llama_free_model(model); + llama_backend_free(); +} diff --git a/retrieval b/retrieval deleted file mode 100755 index dd31789f8..000000000 Binary files a/retrieval and /dev/null differ