Add comments regarding layers output usage.

This commit is contained in:
Andrey Mikhaylov 2024-04-12 11:29:29 +00:00
parent 03284d752e
commit 2c5706f159
2 changed files with 15 additions and 1 deletions

6
.vscode/settings.json vendored Normal file
View File

@ -0,0 +1,6 @@
{
"cmake.configureOnOpen": false,
"files.associations": {
"array": "cpp"
}
}

View File

@ -32,7 +32,13 @@ namespace gcpp {
using GemmaWeightT = GEMMA_WEIGHT_T;
using EmbedderInputT = hwy::bfloat16_t;
using LayersOutputT = std::function<void(int, std::string, const float*, size_t)>;
// Will be called for layers output with:
// - position in the tokens sequence
// - name of the data, p.ex. "tokens", "block.1", "final_norm"
// - ponter to the data array
// - size of the data array
using LayersOutputT =
std::function<void(int, std::string, const float*, size_t)>;
constexpr size_t kPrefillBatchSize = 16;
constexpr bool kSystemPrompt = false;
@ -93,6 +99,8 @@ KVCache CreateKVCache(size_t size_cache_pos, size_t seq_len,
using StreamFunc = std::function<bool(int, float)>;
using AcceptFunc = std::function<bool(int)>;
// layers_output is optional; if set - it will be called with the activations
// output after applying each layer.
void GenerateGemma(Gemma& gemma, size_t max_tokens, size_t max_generated_tokens,
float temperature, const std::vector<int>& prompt,
size_t start_pos, KVCache& kv_cache, hwy::ThreadPool& pool,