From 4ef3da733a05b9e8de0bbae15c633010315a2747 Mon Sep 17 00:00:00 2001
From: Andrey Mikhaylov <andrey.mikhaylov.86@gmail.com>
Date: Fri, 12 Apr 2024 14:54:34 +0000
Subject: [PATCH] Fixed minor things and added comments.

---
 BUILD.bazel     | 2 +-
 debug_prompt.cc | 6 ++----
 gemma/gemma.cc  | 4 ++--
 gemma/gemma.h   | 2 +-
 4 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/BUILD.bazel b/BUILD.bazel
index c1ca793..a504f2f 100644
--- a/BUILD.bazel
+++ b/BUILD.bazel
@@ -176,4 +176,4 @@ cc_binary(
         "@hwy//:thread_pool",
         "@nlohmann_json//:json",
     ],
-)
\ No newline at end of file
+)
diff --git a/debug_prompt.cc b/debug_prompt.cc
index ae4f52d..dd962af 100644
--- a/debug_prompt.cc
+++ b/debug_prompt.cc
@@ -5,9 +5,7 @@
 
 #include "gemma/gemma.h"
 #include "nlohmann/json.hpp"
-// copybara:import_next_line:gemma_cpp
 #include "util/app.h"
-// copybara:import_next_line:gemma_cpp
 #include "util/args.h"
 
 using json = nlohmann::json;
@@ -67,7 +65,7 @@ class OutputJsonLogger {
   json json_output;
 
   gcpp::LayersOutputT layers_output_log_f =
-      [this](int pos, std::string key, const float* values, size_t values_len) {
+      [this](int pos, const std::string& key, const float* values, size_t values_len) {
         std::vector<float> v{values, values + values_len};
         json_output[std::to_string(pos)][key] = v;
       };
@@ -132,4 +130,4 @@ int main(int argc, char** argv) {
   }
 
   return EXIT_SUCCESS;
-}
\ No newline at end of file
+}
diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index 9cfaf1b..be3fca4 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -1180,7 +1180,7 @@ float ComputeCrossEntropyImpl(GemmaImpl<TConfig>& gemma, size_t max_tokens,
              total_entropy / std::log(2.0) / (pos + 1));
     }
     Transformer(token, pos, weights, activations, kv_cache, pool, inner_pool,
-                nullptr);
+                /*layers_output=*/nullptr);
     MatVec<kVocabSize, kModelDim>(weights.embedder_input_embedding, 0,
                                   activations.x.data(),
                                   activations.logits.data(), pool);
@@ -1609,7 +1609,7 @@ void GenerateGemma(Gemma& gemma, RuntimeConfig runtime_config,
       gemma, runtime_config.max_tokens, runtime_config.max_generated_tokens,
       runtime_config.temperature, prompt, start_pos, kv_cache, pool, inner_pool,
       stream_token, [](int) { return true; }, gen, runtime_config.verbosity,
-      nullptr);
+      /*layers_output=*/nullptr);
 }
 
 void CompressWeights(gcpp::Model model, const Path& weights,
diff --git a/gemma/gemma.h b/gemma/gemma.h
index 4120bfa..da37488 100644
--- a/gemma/gemma.h
+++ b/gemma/gemma.h
@@ -38,7 +38,7 @@ using EmbedderInputT = hwy::bfloat16_t;
 // - ponter to the data array
 // - size of the data array
 using LayersOutputT =
-    std::function<void(int, std::string, const float*, size_t)>;
+    std::function<void(int, const std::string&, const float*, size_t)>;
 constexpr size_t kPrefillBatchSize = 16;
 constexpr bool kSystemPrompt = false;