From 6c0388e0495083086f724be98965a016807ae536 Mon Sep 17 00:00:00 2001
From: austinvhuang <austinh@alum.mit.edu>
Date: Thu, 7 Mar 2024 01:14:07 -0500
Subject: [PATCH] [WIP] refine Runtime struct definition

---
 gemma.h | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/gemma.h b/gemma.h
index 3528b50..f6361e1 100644
--- a/gemma.h
+++ b/gemma.h
@@ -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;