Skip to content

Commit 84c0a92

Browse files
Dobiasdclaude
andcommitted
Add minimal GPT-2 chat example
Provides the pieces needed to load keras_hub's gpt2_base_en into a small C++ chat REPL on top of frugally-deep: - keras_export/save_gpt2_backbone_for_fdeep.py: rebuilds the keras_hub GPT2Backbone (or GPT2CausalLM with --with-lm-head) as a plain Keras Functional model using only fdeep-supported primitives, with weights copied over. Verified to match keras_hub within float32 noise on logits. - keras_export/save_gpt2_weights_bin.py: dumps the rebuilt model's weights to a flat binary file consumed by the cached inference engine. - include/fdeep/llm/gpt2_bpe.hpp: byte-level BPE tokenizer (vocab.json + merges.txt). ASCII-focused pre-tokenization; round-trip safe and matches keras_hub on the test set. - include/fdeep/llm/gpt2_generator.hpp: slow but correct generation loop on top of fdeep::model::predict, useful as a reference path. - include/fdeep/llm/gpt2_cached.hpp: stateful Eigen-based GPT-2 forward with a per-layer K/V cache. Loads weights from the binary file; prefill() seeds the cache, step() advances by one token in roughly constant time. - examples/gpt2_chat: simple stdin REPL using the cached engine. Greedy / temperature / top-k sampling. ~18 ms/token sustained on CPU for gpt2_base_en at seq_len 256. The cached engine bypasses fdeep::model::predict because frugally-deep's runtime is stateless and a per-step KV cache cleanly fits a custom forward path. The slow generator path remains available for users who'd rather stay on the generic runtime. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 0abe158 commit 84c0a92

7 files changed

Lines changed: 1615 additions & 0 deletions

File tree

examples/gpt2_chat/CMakeLists.txt

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
# Example: minimal GPT-2 chat demo. Disabled by default; enable with
2+
# cmake -DFDEEP_BUILD_GPT2_CHAT=ON ..
3+
# The example needs the rebuilt model JSON and the GPT-2 vocab/merges, see
4+
# keras_export/save_gpt2_backbone_for_fdeep.py and the README in this directory.
5+
6+
option(FDEEP_BUILD_GPT2_CHAT "Build the GPT-2 chat demo example" OFF)
7+
8+
if(FDEEP_BUILD_GPT2_CHAT)
9+
add_executable(gpt2_chat main.cpp)
10+
target_link_libraries(gpt2_chat PRIVATE fdeep)
11+
if(NOT MSVC)
12+
target_compile_options(gpt2_chat PRIVATE -O3)
13+
endif()
14+
endif()

examples/gpt2_chat/main.cpp

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
// Minimal GPT-2 chat demo on top of frugally-deep.
2+
//
3+
// Usage:
4+
// gpt2_chat <weights.bin> <vocab.json> <merges.txt> [max_new_tokens] [temperature] [top_k]
5+
//
6+
// Reads a prompt per line from stdin and prints the continuation. A blank
7+
// line exits. Uses a stateful key/value cache (see
8+
// include/fdeep/llm/gpt2_cached.hpp), so each new token costs roughly
9+
// constant time rather than re-encoding the whole prefix.
10+
//
11+
// The weights binary is produced by
12+
// keras_export/save_gpt2_weights_bin.py --output weights.bin
13+
14+
#include <fdeep/llm/gpt2_bpe.hpp>
15+
#include <fdeep/llm/gpt2_cached.hpp>
16+
17+
#include <algorithm>
18+
#include <chrono>
19+
#include <cmath>
20+
#include <cstdint>
21+
#include <iostream>
22+
#include <random>
23+
#include <string>
24+
#include <vector>
25+
26+
namespace {
27+
28+
int sample_logits(const std::vector<float>& logits, float temperature,
29+
std::size_t top_k, std::mt19937_64& rng)
30+
{
31+
if (temperature <= 0.0f) {
32+
std::size_t best = 0;
33+
float best_val = logits[0];
34+
for (std::size_t i = 1; i < logits.size(); ++i) {
35+
if (logits[i] > best_val) {
36+
best_val = logits[i];
37+
best = i;
38+
}
39+
}
40+
return static_cast<int>(best);
41+
}
42+
43+
std::vector<std::pair<float, int>> cand;
44+
cand.reserve(logits.size());
45+
const float inv_t = 1.0f / temperature;
46+
for (std::size_t i = 0; i < logits.size(); ++i) {
47+
cand.emplace_back(logits[i] * inv_t, static_cast<int>(i));
48+
}
49+
if (top_k > 0 && top_k < cand.size()) {
50+
std::partial_sort(cand.begin(),
51+
cand.begin() + static_cast<std::ptrdiff_t>(top_k), cand.end(),
52+
[](const auto& a, const auto& b) { return a.first > b.first; });
53+
cand.resize(top_k);
54+
}
55+
float m = cand[0].first;
56+
for (const auto& c : cand) m = std::max(m, c.first);
57+
float total = 0.0f;
58+
for (auto& c : cand) {
59+
c.first = std::exp(c.first - m);
60+
total += c.first;
61+
}
62+
std::uniform_real_distribution<float> dist(0.0f, total);
63+
const float pick = dist(rng);
64+
float acc = 0.0f;
65+
for (const auto& c : cand) {
66+
acc += c.first;
67+
if (acc >= pick) return c.second;
68+
}
69+
return cand.back().second;
70+
}
71+
72+
} // namespace
73+
74+
int main(int argc, char** argv)
75+
{
76+
if (argc < 4) {
77+
std::cerr << "usage: " << argv[0]
78+
<< " <weights.bin> <vocab.json> <merges.txt>"
79+
" [max_new_tokens] [temperature] [top_k] [max_seq_len]\n";
80+
return 1;
81+
}
82+
const std::string weights_path = argv[1];
83+
const std::string vocab_path = argv[2];
84+
const std::string merges_path = argv[3];
85+
const std::size_t max_new = (argc > 4) ? std::stoul(argv[4]) : 100;
86+
const float temperature = (argc > 5) ? std::stof(argv[5]) : 0.0f;
87+
const std::size_t top_k = (argc > 6) ? std::stoul(argv[6]) : 0;
88+
const std::size_t max_seq_len = (argc > 7) ? std::stoul(argv[7]) : 256;
89+
90+
std::cerr << "loading tokenizer...\n";
91+
fdeep::llm::gpt2_bpe_tokenizer tok(vocab_path, merges_path);
92+
93+
std::cerr << "loading weights...\n";
94+
const auto t0 = std::chrono::steady_clock::now();
95+
fdeep::llm::gpt2_cached_model gpt(weights_path, max_seq_len);
96+
const auto t1 = std::chrono::steady_clock::now();
97+
std::cerr << "loaded in "
98+
<< std::chrono::duration<double>(t1 - t0).count() << " s\n";
99+
100+
std::mt19937_64 rng(static_cast<std::uint64_t>(
101+
std::chrono::steady_clock::now().time_since_epoch().count()));
102+
103+
std::cerr << "ready (max_seq_len=" << max_seq_len
104+
<< ", temperature=" << temperature
105+
<< ", top_k=" << top_k
106+
<< "). type a prompt and press enter; blank line quits.\n";
107+
108+
std::string line;
109+
while (true) {
110+
std::cout << "> " << std::flush;
111+
if (!std::getline(std::cin, line)) break;
112+
if (line.empty()) break;
113+
114+
gpt.reset();
115+
const auto prompt_ids = tok.encode(line);
116+
if (prompt_ids.empty()) continue;
117+
if (prompt_ids.size() >= max_seq_len) {
118+
std::cerr << "[prompt is " << prompt_ids.size()
119+
<< " tokens; max_seq_len is " << max_seq_len
120+
<< "]\n";
121+
continue;
122+
}
123+
124+
std::cout << line;
125+
std::cout.flush();
126+
127+
const auto t_pre0 = std::chrono::steady_clock::now();
128+
auto logits = gpt.prefill(prompt_ids);
129+
const auto t_pre1 = std::chrono::steady_clock::now();
130+
131+
std::vector<int> generated;
132+
const auto t_gen0 = std::chrono::steady_clock::now();
133+
for (std::size_t step = 0; step < max_new; ++step) {
134+
const int next_id = sample_logits(logits, temperature, top_k, rng);
135+
generated.push_back(next_id);
136+
std::cout << tok.decode({ next_id }) << std::flush;
137+
if (next_id == tok.eos_token_id()) break;
138+
if (gpt.cur_len() >= gpt.max_seq_len()) break;
139+
logits = gpt.step(next_id);
140+
}
141+
const auto t_gen1 = std::chrono::steady_clock::now();
142+
std::cout << "\n["
143+
<< "prefill " << prompt_ids.size() << " tok "
144+
<< std::chrono::duration<double>(t_pre1 - t_pre0).count() << " s, "
145+
<< "decode " << generated.size() << " tok "
146+
<< std::chrono::duration<double>(t_gen1 - t_gen0).count() << " s, "
147+
<< (generated.empty() ? 0.0 :
148+
1000.0 * std::chrono::duration<double>(t_gen1 - t_gen0).count()
149+
/ static_cast<double>(generated.size()))
150+
<< " ms/tok]\n";
151+
}
152+
return 0;
153+
}

0 commit comments

Comments
 (0)