[WIP] refine Runtime struct definition

This commit is contained in:
austinvhuang 2024-03-07 01:14:07 -05:00
parent 49e654258d
commit 6c0388e049
1 changed files with 6 additions and 5 deletions

11
gemma.h
View File

@ -64,13 +64,14 @@ struct KVCache {
enum class Model { GEMMA_2B, GEMMA_7B };
enum class ModelTraining { GEMMA_IT, GEMMA_PT };
// TODO: incorporate
struct InferenceParams {
Model model;
// TODO: Incorporate this
struct Runtime {
Model model_type;
ModelTraining model_training;
size_t max_generated_tokens;
size_t max_tokens;
size_t max_generated_tokens;
float temperature;
std::mt19937 gen;
};
struct LoaderArgs : public ArgsBase<LoaderArgs> {
@ -212,7 +213,7 @@ void GenerateGemma(Gemma& gemma, size_t max_tokens, size_t max_generated_tokens,
float temperature, const std::vector<int>& prompt,
size_t start_pos, hwy::ThreadPool& pool,
hwy::ThreadPool& inner_pool, const StreamFunc& stream_token,
const AcceptFunc& accept_token, std::mt19937& g,
const AcceptFunc& accept_token, std::mt19937& gen,
int verbosity);
constexpr int EOS_ID = 1;