Skip to content

Commit 2bc0bfb

Browse files
JohannesGaesslerJohannesGaessler
committed
lookup: use static data from text corpus
1 parent 8f62e5d commit 2bc0bfb

File tree

4 files changed

+220
-6
lines changed

4 files changed

+220
-6
lines changed

‎Makefile

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,8 @@ lookahead: examples/lookahead/lookahead.cpp ggml.o llama.o $(COMMON_DEPS) $(OBJS
784784
lookup: examples/lookup/lookup.cpp ggml.o llama.o $(COMMON_DEPS) $(OBJS)
785785
$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)
786786
$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, $<) -o $@ $(LDFLAGS)
787+
$(CXX) $(CXXFLAGS) -c examples/lookup/lookup-create.cpp -o $(call GET_OBJ_FILE, examples/lookup/lookup-create.cpp)
788+
$(CXX) $(CXXFLAGS) $(filter-out %.h $<,$^) $(call GET_OBJ_FILE, examples/lookup/lookup-create.cpp) -o lookup-create $(LDFLAGS)
787789

788790
passkey: examples/passkey/passkey.cpp ggml.o llama.o $(COMMON_DEPS) $(OBJS)
789791
$(CXX) $(CXXFLAGS) -c $< -o $(call GET_OBJ_FILE, $<)

‎examples/lookup/lookup-create.cpp

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
#include "common.h"
2+
#include "ggml.h"
3+
#include "llama.h"
4+
5+
#include <cstdint>
6+
#include <fstream>
7+
#include <iostream>
8+
#include <string>
9+
#include <unordered_map>
10+
#include <vector>
11+
12+
typedef std::unordered_map<llama_token, int32_t> token_hashmap; // token -> number of times token has been seen
13+
typedef std::unordered_map<uint64_t, token_hashmap> all_token_hashmap; // n-gram -> empirical distribution of following tokens
14+
constexpr int ngram_size = 2;
15+
16+
int main(int argc, char ** argv){
17+
gpt_params params;
18+
19+
if (!gpt_params_parse(argc, argv, params)) {
20+
return 1;
21+
}
22+
// init llama.cpp
23+
llama_backend_init(params.numa);
24+
25+
llama_model * model = NULL;
26+
llama_context * ctx = NULL;
27+
28+
// load the model
29+
std::tie(model, ctx) = llama_init_from_gpt_params(params);
30+
GGML_ASSERT(model != nullptr);
31+
32+
// tokenize the prompt
33+
const bool add_bos = llama_should_add_bos_token(model);
34+
35+
const char * static_input_file = "./wikitext-2-raw/wiki.train.raw";
36+
std::ifstream file(static_input_file);
37+
if (!file) {
38+
fprintf(stderr, "error: failed to open file '%s'\n", static_input_file);
39+
exit(1);
40+
}
41+
std::string static_input;
42+
std::copy(std::istreambuf_iterator<char>(file), std::istreambuf_iterator<char>(), back_inserter(static_input));
43+
if (!static_input.empty() && static_input.back() == '\n') {
44+
static_input.pop_back();
45+
}
46+
std::vector<llama_token> inp_static;
47+
inp_static = ::llama_tokenize(ctx, static_input, add_bos, true);
48+
fprintf(stderr, "lookup-create: tokenization done\n");
49+
50+
auto update_hashmaps = [](all_token_hashmap * atc, const llama_token * inp_data, const int inp_size, const int nnew) -> void {
51+
// atcs = all_token_counts: the hashmaps to modify.
52+
// inp_data: the token sequence on which the hashmaps are based.
53+
// inp_size: the current size of inp_data.
54+
// nnew: how many new tokens have been appended to inp_data since the last call to this function.
55+
//
56+
// In order to get correct results inp_data can ONLY BE APPENDED TO.
57+
// Changes in the middle need a complete rebuild.
58+
59+
const int i_start = std::max(inp_size - nnew, ngram_size);
60+
const int64_t t_start_ms = ggml_time_ms();
61+
int percentage_done = 0;
62+
for (int i = i_start; i < inp_size; ++i) {
63+
const int ngram_start = i - ngram_size;
64+
uint64_t ngram = inp_data[ngram_start];
65+
for (int j = ngram_start; j < ngram_start + ngram_size; ++j) {
66+
const uint64_t ngram_part = inp_data[j];
67+
ngram <<= 16;
68+
ngram |= ngram_part;
69+
}
70+
const llama_token token = inp_data[i];
71+
72+
all_token_hashmap::iterator token_counts_it = atc->find(ngram);
73+
if (token_counts_it == atc->end()) {
74+
token_hashmap token_counts;
75+
token_counts.emplace(token, 1);
76+
atc->emplace(ngram, token_counts);
77+
} else {
78+
token_hashmap::iterator tc_it = token_counts_it->second.find(token);
79+
if (tc_it == token_counts_it->second.end()) {
80+
token_counts_it->second.emplace(token, 1);
81+
} else {
82+
tc_it->second++;
83+
}
84+
}
85+
86+
if (i >= inp_size*(percentage_done + 1)/100) {
87+
++percentage_done;
88+
89+
const int64_t t_now_ms = ggml_time_ms();
90+
const int64_t eta_ms = (100 - percentage_done) * (t_now_ms - t_start_ms) / percentage_done;
91+
const int64_t eta_min = eta_ms / (60*1000);
92+
const int64_t eta_s = (eta_ms - eta_min) / 1000;
93+
94+
fprintf(stderr, "lookup-create: %02d%% done, ETA: %02ld:%02ld\n", percentage_done, eta_min, eta_s);
95+
}
96+
}
97+
};
98+
99+
all_token_hashmap atc;
100+
update_hashmaps(&atc, inp_static.data(), inp_static.size(), inp_static.size());
101+
102+
std::ofstream file_out("lookup.bin", std::ios::binary);
103+
for (std::pair<uint64_t, token_hashmap> item : atc) {
104+
const uint64_t ngram = item.first;
105+
token_hashmap token_counts = item.second;
106+
GGML_ASSERT(!token_counts.empty());
107+
const int32_t ntokens = token_counts.size();
108+
109+
110+
file_out.write(reinterpret_cast<const char *>(&ngram), sizeof(uint64_t));
111+
file_out.write(reinterpret_cast<const char *>(&ntokens), sizeof(int32_t));
112+
for (std::pair<llama_token, int32_t> item2 : token_counts) {
113+
const llama_token token = item2.first;
114+
const int32_t count = item2.second;
115+
file_out.write(reinterpret_cast<const char *>(&token), sizeof(llama_token));
116+
file_out.write(reinterpret_cast<const char *>(&count), sizeof(int32_t));
117+
}
118+
}
119+
}

‎examples/lookup/lookup.cpp

Lines changed: 89 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,13 @@
55
#include <cmath>
66
#include <cstdint>
77
#include <cstdio>
8+
#include <fstream>
89
#include <string>
910
#include <vector>
1011
#include <unordered_map>
1112

1213
// Data structures to map n-grams to empirical token probabilities:
13-
typedef std::unordered_map<llama_token, int> token_hashmap; // token -> number of times token has been seen
14+
typedef std::unordered_map<llama_token, int32_t> token_hashmap; // token -> number of times token has been seen
1415
typedef std::unordered_map<uint64_t, token_hashmap> all_token_hashmap; // n-gram -> empirical distribution of following tokens
1516
// n-grams are encoded as 64 bit integers with each of the 4 16 bit sections representing a token id.
1617
// This way no custom hashing function for the n-grams is needed.
@@ -22,7 +23,7 @@ static_assert(ngram_max <= sizeof(uint64_t)/2, "A 64 bit integer can only hold i
2223

2324
// If sample size or percentage in context are below these thresholds the draft is aborted early:
2425
constexpr float draft_min_sample_size[ngram_max] = { 2, 2, 1, 1};
25-
constexpr float draft_min_percent[ngram_max] = {66, 50, 50, 50};
26+
constexpr float draft_min_percent[ngram_max] = {50, 50, 50, 50};
2627

2728
int main(int argc, char ** argv){
2829
gpt_params params;
@@ -101,12 +102,43 @@ int main(int argc, char ** argv){
101102
};
102103

103104
all_token_hashmap all_token_counts[ngram_max-ngram_min+1];
105+
all_token_hashmap static_all_token_counts;
104106
int64_t t_draft_us = 0;
105107

106108
{
107109
// Fill up hashmaps with tokens from user input:
108110
const int64_t t_start_draft_us = ggml_time_us();
109111
update_hashmaps(all_token_counts, inp.data(), inp.size(), inp.size());
112+
113+
const char * hashmap_file_name = "lookup.bin";
114+
std::ifstream hashmap_file(hashmap_file_name, std::ios::binary);
115+
if (!hashmap_file) {
116+
fprintf(stderr, "error: failed to open file '%s'\n", hashmap_file_name);
117+
exit(1);
118+
}
119+
uint64_t ngram;
120+
int32_t ntokens;
121+
llama_token token;
122+
int32_t count;
123+
124+
char * ngramc = reinterpret_cast<char*>(&ngram);
125+
char * ntokensc = reinterpret_cast<char*>(&ntokens);
126+
char * tokenc = reinterpret_cast<char*>(&token);
127+
char * countc = reinterpret_cast<char*>(&count);
128+
while(hashmap_file.read(ngramc, sizeof(uint64_t))) {
129+
GGML_ASSERT(hashmap_file.read(ntokensc, sizeof(int32_t)));
130+
token_hashmap token_counts;
131+
132+
for (int i = 0; i < ntokens; ++i) {
133+
GGML_ASSERT(hashmap_file.read(tokenc, sizeof(llama_token)));
134+
GGML_ASSERT(hashmap_file.read(countc, sizeof(int32_t)));
135+
token_counts.emplace(token, count);
136+
}
137+
138+
static_all_token_counts.emplace(ngram, token_counts);
139+
}
140+
GGML_ASSERT(hashmap_file.eof());
141+
110142
t_draft_us += ggml_time_us() - t_start_draft_us;
111143
}
112144

@@ -249,6 +281,20 @@ int main(int argc, char ** argv){
249281

250282
while ((int) draft.size()-1 < n_draft) {
251283
bool draft_success = false;
284+
285+
const int static_ngram_start = inp_size-2 + draft.size()-1;
286+
uint64_t static_ngram = get_token(inp, draft, static_ngram_start);
287+
for (int j = static_ngram_start; j < static_ngram_start + 2; ++j) {
288+
const uint64_t ngram_part = get_token(inp, draft, j);
289+
static_ngram <<= 16;
290+
static_ngram |= ngram_part;
291+
}
292+
all_token_hashmap::iterator static_token_counts_it = static_all_token_counts.find(static_ngram);
293+
token_hashmap static_token_counts;
294+
if (static_token_counts_it != static_all_token_counts.end()) {
295+
static_token_counts = static_token_counts_it->second;
296+
}
297+
252298
for (int ngram_size = ngram_max; ngram_size >= ngram_min; --ngram_size) {
253299
if (ngram_size > inp_size) {
254300
continue;
@@ -271,16 +317,21 @@ int main(int argc, char ** argv){
271317
const token_hashmap token_counts = token_counts_it->second;
272318

273319
int max_count = 0;
320+
int max_count_static = 0;
274321
int sum_count = 0;
275322
llama_token max_token = -1;
276323

277324
for (std::pair<llama_token, int> tc : token_counts) {
278325
const llama_token token = tc.first;
279-
const llama_token count = tc.second;
280326

281-
if (count > max_count) {
282-
max_token = token;
283-
max_count = count;
327+
token_hashmap::iterator stc_it = static_token_counts.find(token);
328+
const int32_t count = tc.second;
329+
const int32_t count_static = stc_it != static_token_counts.end() ? 100*stc_it->second : 1;
330+
331+
if (count*count_static > max_count*max_count_static) {
332+
max_token = token;
333+
max_count = count;
334+
max_count_static = count_static;
284335
}
285336
sum_count += count;
286337
}
@@ -300,6 +351,38 @@ int main(int argc, char ** argv){
300351
break;
301352
}
302353

354+
if (!draft_success) {
355+
int max_count = 0;
356+
int sum_count = 0;
357+
llama_token max_token = -1;
358+
359+
for (std::pair<llama_token, int> tc : static_token_counts) {
360+
const llama_token token = tc.first;
361+
const int32_t count = tc.second;
362+
363+
if (count > max_count) {
364+
max_token = token;
365+
max_count = count;
366+
}
367+
sum_count += count;
368+
}
369+
370+
// Skip this candidate if the sample size is too low:
371+
if (sum_count < draft_min_sample_size[2-1]) {
372+
break;
373+
}
374+
// skip this candidate if the empirically most likely token following this token is not likely enough:
375+
if (100*max_count < draft_min_percent[2-1]*sum_count) {
376+
break;
377+
}
378+
379+
LOG(" - draft candidate: token=%d count=%d\n", max_token, max_count);
380+
llama_batch_add(batch_tgt, max_token, n_past + draft.size(), { 0 }, true);
381+
draft.push_back(max_token);
382+
draft_success = true;
383+
break;
384+
}
385+
303386
if (!draft_success) {
304387
break;
305388
}

‎scripts/get-wikitext-103.sh

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
#!/bin/bash
2+
3+
wget https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip
4+
5+
echo "Usage:"
6+
echo ""
7+
echo " ./perplexity -m model.gguf -f wiki.test.raw [other params]"
8+
echo ""
9+
10+
exit 0

0 commit comments

Comments
 (0)