Skip to content

Commit eb760a9

Browse files
mscheong01mscheong01
committed
add retrieval example
1 parent c5b8595 commit eb760a9

File tree

7 files changed

+335
-1
lines changed

7 files changed

+335
-1
lines changed

‎.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ models-mnt
7272
/batched-bench
7373
/export-lora
7474
/finetune
75+
/retrieval
7576
/speculative
7677
/parallel
7778
/train-text-from-scratch

‎Makefile

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
BUILD_TARGETS = \
33
main quantize quantize-stats perplexity imatrix embedding vdot q8dot train-text-from-scratch convert-llama2c-to-ggml \
44
simple batched batched-bench save-load-state server gguf llama-bench libllava.a llava-cli baby-llama beam-search \
5-
speculative infill tokenize benchmark-matmult parallel finetune export-lora lookahead lookup passkey gritlm tests/test-c.o
5+
retrieval speculative infill tokenize benchmark-matmult parallel finetune export-lora lookahead lookup passkey gritlm tests/test-c.o
66

77
# Binaries only useful for tests
88
TEST_TARGETS = \
@@ -794,6 +794,10 @@ export-lora: examples/export-lora/export-lora.cpp ggml.o common/common.h $(OBJS)
794794
$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
795795
$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)
796796

797+
retrieval: examples/retrieval/retrieval.cpp ggml.o llama.o $(COMMON_DEPS) $(OBJS)
798+
$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
799+
$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)
800+
797801
speculative: examples/speculative/speculative.cpp ggml.o llama.o $(COMMON_DEPS) grammar-parser.o $(OBJS)
798802
$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
799803
$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)

‎common/common.cpp

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,43 @@ static bool gpt_params_find_arg(int argc, char ** argv, gpt_params & params, int
276276
}
277277
return true;
278278
}
279+
if (arg == "--context-files") {
280+
if (++i >= argc) {
281+
invalid_param = true;
282+
return true;
283+
}
284+
while(true) {
285+
std::ifstream file(argv[i]);
286+
if (!file) {
287+
fprintf(stderr, "error: failed to open file '%s'\n", argv[i]);
288+
invalid_param = true;
289+
break;
290+
}
291+
// store the external file name in params
292+
params.context_files.push_back(argv[i]);
293+
if (i + 1 >= argc || argv[i + 1][0] == '-') {
294+
break;
295+
}
296+
i++;
297+
}
298+
return true;
299+
}
300+
if (arg == "--chunk-size") {
301+
if (++i >= argc) {
302+
invalid_param = true;
303+
return true;
304+
}
305+
params.chunk_size = std::stoi(argv[i]);
306+
return true;
307+
}
308+
if (arg == "--chunk-separator") {
309+
if (++i >= argc) {
310+
invalid_param = true;
311+
return true;
312+
}
313+
params.chunk_separator = argv[i];
314+
return true;
315+
}
279316
if (arg == "-n" || arg == "--n-predict") {
280317
if (++i >= argc) {
281318
invalid_param = true;
@@ -1282,6 +1319,11 @@ void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
12821319
printf(" prompt file to start generation.\n");
12831320
printf(" -bf FNAME, --binary-file FNAME\n");
12841321
printf(" binary file containing multiple choice tasks.\n");
1322+
printf(" --context-files FNAME1 FNAME2...\n");
1323+
printf(" files containing context to embed.\n");
1324+
printf(" --chunk-size N minimum length of embedded text chunk (default:%d)\n", params.chunk_size);
1325+
printf(" --chunk-separator STRING\n");
1326+
printf(" string to separate chunks (default: newline)\n");
12851327
printf(" -n N, --n-predict N number of tokens to predict (default: %d, -1 = infinity, -2 = until context filled)\n", params.n_predict);
12861328
printf(" -c N, --ctx-size N size of the prompt context (default: %d, 0 = loaded from model)\n", params.n_ctx);
12871329
printf(" -b N, --batch-size N logical maximum batch size (default: %d)\n", params.n_batch);

‎common/common.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ struct gpt_params {
7979
float yarn_beta_slow = 1.0f; // YaRN high correction dim
8080
int32_t yarn_orig_ctx = 0; // YaRN original context length
8181
float defrag_thold = -1.0f; // KV cache defragmentation threshold
82+
std::vector<std::string> context_files = {}; // context files to embed
83+
int32_t chunk_size = 64; // chunk size for context embedding
84+
std::string chunk_separator = "\n"; // chunk separator for context embedding
8285

8386
ggml_numa_strategy numa = GGML_NUMA_STRATEGY_DISABLED;
8487

‎examples/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ else()
3434
add_subdirectory(perplexity)
3535
add_subdirectory(quantize)
3636
add_subdirectory(quantize-stats)
37+
add_subdirectory(retrieval)
3738
add_subdirectory(save-load-state)
3839
add_subdirectory(simple)
3940
add_subdirectory(passkey)

‎examples/retrieval/CMakeLists.txt

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
set(TARGET retrieval)
2+
add_executable(${TARGET} retrieval.cpp)
3+
install(TARGETS ${TARGET} RUNTIME)
4+
target_link_libraries(${TARGET} PRIVATE common llama ${CMAKE_THREAD_LIBS_INIT})
5+
target_compile_features(${TARGET} PRIVATE cxx_std_11)

‎examples/retrieval/retrieval.cpp

Lines changed: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
1+
#include "common.h"
2+
#include "llama.h"
3+
4+
#include <algorithm>
5+
#include <fstream>
6+
7+
struct chunk {
8+
// filename
9+
std::string filename;
10+
// original file position
11+
int64_t filepos;
12+
// original text data
13+
std::string textdata = "";
14+
// tokenized text data
15+
std::vector<std::int32_t> tokens;
16+
// embedding
17+
std::vector<float> embedding;
18+
// cosin similarity
19+
float similarity;
20+
};
21+
22+
// chunk file data to chunks of size >= chunk_size
23+
// chunk_separator is the separator between chunks
24+
static std::vector<chunk> chunk_file(const std::string filename, int chunk_size, std::string chunk_separator) {
25+
std::vector<chunk> chunks;
26+
std::ifstream f(filename.c_str());
27+
28+
if (!f.is_open()) {
29+
fprintf(stderr, "Error: could not open file %s\n", filename.c_str());
30+
return chunks;
31+
}
32+
33+
chunk current_chunk;
34+
char buffer[chunk_size];
35+
int64_t filepos = 0;
36+
std::string current = "";
37+
while (f.read(buffer, chunk_size)) {
38+
current += std::string(buffer, f.gcount());
39+
size_t pos;
40+
while ((pos = current.find(chunk_separator)) != std::string::npos) {
41+
current_chunk.textdata += current.substr(0, pos + chunk_separator.size());
42+
if ((int) current_chunk.textdata.size() > chunk_size) {
43+
// save chunk
44+
current_chunk.filepos = filepos;
45+
current_chunk.filename = filename;
46+
chunks.push_back(current_chunk);
47+
// update filepos
48+
filepos += (int) current_chunk.textdata.size();
49+
// reset current_chunk
50+
current_chunk = chunk();
51+
}
52+
current = current.substr(pos + chunk_separator.size());
53+
}
54+
55+
}
56+
// add leftover data to last chunk
57+
if (current_chunk.textdata.size() > 0) {
58+
if (chunks.empty()) {
59+
current_chunk.filepos = filepos;
60+
current_chunk.filename = filename;
61+
chunks.push_back(current_chunk);
62+
} else {
63+
chunks.back().textdata += current_chunk.textdata;
64+
}
65+
}
66+
f.close();
67+
return chunks;
68+
}
69+
70+
static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, int seq_id) {
71+
for (size_t i = 0; i < tokens.size(); i++) {
72+
llama_batch_add(batch, tokens[i], i, { seq_id }, i == tokens.size() - 1);
73+
}
74+
}
75+
76+
static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) {
77+
// clear previous kv_cache values (irrelevant for embeddings)
78+
llama_kv_cache_clear(ctx);
79+
80+
// run model
81+
fprintf(stderr, "%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);
82+
if (llama_decode(ctx, batch) < 0) {
83+
fprintf(stderr, "%s : failed to decode\n", __func__);
84+
}
85+
86+
for (int i = 0; i < batch.n_tokens; i++) {
87+
if (!batch.logits[i]) {
88+
continue;
89+
}
90+
91+
// try to get sequence embeddings - supported only when pooling_type is not NONE
92+
const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);
93+
if (embd == NULL) {
94+
embd = llama_get_embeddings_ith(ctx, i);
95+
if (embd == NULL) {
96+
fprintf(stderr, "%s: failed to get embeddings for token %d\n", __func__, i);
97+
continue;
98+
}
99+
}
100+
101+
float * out = output + batch.seq_id[i][0] * n_embd;
102+
llama_embd_normalize(embd, out, n_embd);
103+
}
104+
}
105+
106+
int main(int argc, char ** argv) {
107+
gpt_params params;
108+
109+
if (!gpt_params_parse(argc, argv, params)) {
110+
return 1;
111+
}
112+
113+
if (params.chunk_size <= 0) {
114+
fprintf(stderr, "chunk_size must be positive\n");
115+
return 1;
116+
}
117+
if (params.context_files.empty()) {
118+
fprintf(stderr, "context_files must be specified\n");
119+
return 1;
120+
}
121+
params.embedding = true;
122+
123+
print_build_info();
124+
125+
if (params.seed == LLAMA_DEFAULT_SEED) {
126+
params.seed = time(NULL);
127+
}
128+
129+
printf("processing files:\n");
130+
for (auto & context_file : params.context_files) {
131+
printf("%s\n", context_file.c_str());
132+
}
133+
134+
std::vector<chunk> chunks;
135+
for (auto & context_file : params.context_files) {
136+
std::vector<chunk> file_chunk = chunk_file(context_file, params.chunk_size, params.chunk_separator);
137+
chunks.insert(chunks.end(), file_chunk.begin(), file_chunk.end());
138+
}
139+
printf("Number of chunks: %ld\n", chunks.size());
140+
141+
llama_backend_init();
142+
llama_numa_init(params.numa);
143+
144+
llama_model * model;
145+
llama_context * ctx;
146+
147+
// load the model
148+
std::tie(model, ctx) = llama_init_from_gpt_params(params);
149+
if (model == NULL) {
150+
fprintf(stderr, "%s: error: unable to load model\n", __func__);
151+
return 1;
152+
}
153+
154+
const int n_ctx_train = llama_n_ctx_train(model);
155+
const int n_ctx = llama_n_ctx(ctx);
156+
157+
if (n_ctx > n_ctx_train) {
158+
fprintf(stderr, "%s: warning: model was trained on only %d context tokens (%d specified)\n",
159+
__func__, n_ctx_train, n_ctx);
160+
}
161+
162+
// print system information
163+
{
164+
fprintf(stderr, "\n");
165+
fprintf(stderr, "%s\n", get_system_info(params).c_str());
166+
}
167+
168+
// max batch size
169+
const uint64_t n_batch = params.n_batch;
170+
GGML_ASSERT(params.n_batch >= params.n_ctx);
171+
172+
// tokenize the prompts and trim
173+
for (auto & chunk : chunks) {
174+
auto inp = ::llama_tokenize(ctx, chunk.textdata, true, false);
175+
if (inp.size() > n_batch) {
176+
inp.resize(n_batch);
177+
}
178+
// add eos if not present
179+
if (inp.empty() || inp.back() != llama_token_eos(model)) {
180+
inp.push_back(llama_token_eos(model));
181+
}
182+
chunk.tokens = inp;
183+
}
184+
185+
// tokenization stats
186+
if (params.verbose_prompt) {
187+
for (int i = 0; i < (int) chunks.size(); i++) {
188+
fprintf(stderr, "%s: prompt %d: '%s'\n", __func__, i, chunks[i].textdata.c_str());
189+
fprintf(stderr, "%s: number of tokens in prompt = %zu\n", __func__, chunks[i].tokens.size());
190+
for (int j = 0; j < (int) chunks[i].tokens.size(); j++) {
191+
fprintf(stderr, "%6d -> '%s'\n", chunks[i].tokens[j], llama_token_to_piece(ctx, chunks[i].tokens[j]).c_str());
192+
}
193+
fprintf(stderr, "\n\n");
194+
}
195+
}
196+
197+
// initialize batch
198+
const int n_chunks = chunks.size();
199+
struct llama_batch batch = llama_batch_init(n_batch, 0, 1);
200+
201+
// allocate output
202+
const int n_embd = llama_n_embd(model);
203+
std::vector<float> embeddings(n_chunks * n_embd, 0);
204+
float * emb = embeddings.data();
205+
206+
// break into batches
207+
int p = 0; // number of prompts processed already
208+
int s = 0; // number of prompts in current batch
209+
for (int k = 0; k < n_chunks; k++) {
210+
// clamp to n_batch tokens
211+
auto & inp = chunks[k].tokens;
212+
213+
const uint64_t n_toks = inp.size();
214+
215+
// encode if at capacity
216+
if (batch.n_tokens + n_toks > n_batch) {
217+
float * out = emb + p * n_embd;
218+
batch_decode(ctx, batch, out, s, n_embd);
219+
llama_batch_clear(batch);
220+
p += s;
221+
s = 0;
222+
}
223+
224+
// add to batch
225+
batch_add_seq(batch, inp, s);
226+
s += 1;
227+
}
228+
229+
// final batch
230+
float * out = emb + p * n_embd;
231+
batch_decode(ctx, batch, out, s, n_embd);
232+
233+
// save embeddings to chunks
234+
for (int i = 0; i < n_chunks; i++) {
235+
chunks[i].embedding = std::vector<float>(emb + i * n_embd, emb + (i + 1) * n_embd);
236+
}
237+
238+
// start loop, receive query and return top k similar chunks based on cosine similarity
239+
std::string query;
240+
while (true) {
241+
printf("Enter query: ");
242+
std::getline(std::cin, query);
243+
if (query == "exit" || query == "quit" || query == "q") {
244+
break;
245+
}
246+
std::vector<int32_t> query_tokens = llama_tokenize(ctx, query, true);
247+
248+
struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1);
249+
batch_add_seq(query_batch, query_tokens, 0);
250+
float * query_emb = new float[n_embd];
251+
batch_decode(ctx, query_batch, query_emb, 1, n_embd);
252+
std::vector<float> query_embedding(query_emb, query_emb + n_embd);
253+
delete[] query_emb;
254+
llama_batch_clear(query_batch);
255+
256+
for (int i = 0; i < n_chunks; i++) {
257+
float similarity = llama_embd_similarity_cos(chunks[i].embedding.data(), query_embedding.data(), n_embd);
258+
chunks[i].similarity = similarity;
259+
}
260+
std::sort(chunks.begin(), chunks.end(), [](chunk & a, chunk & b) {
261+
return a.similarity > b.similarity;
262+
});
263+
printf("Top %d similar chunks:\n", params.sparams.top_k);
264+
for (int i = 0; i < std::min(params.sparams.top_k, (int) chunks.size()); i++) {
265+
printf("filename: %s\n", chunks[i].filename.c_str());
266+
printf("filepos: %lld\n", chunks[i].filepos);
267+
printf("similarity: %f\n", chunks[i].similarity);
268+
printf("textdata:\n%s\n", chunks[i].textdata.c_str());
269+
printf("--------------------\n");
270+
}
271+
}
272+
273+
// clean up
274+
llama_print_timings(ctx);
275+
llama_free(ctx);
276+
llama_free_model(model);
277+
llama_backend_free();
278+
}

0 commit comments

Comments
 (0)