Use a callback to report TimeStampsAndStats to users

We keep track of num_prompt_tokens and num_generated_tokens along with
the timers. Then we report everything back to user.

Author: kirklandsign

examples/models/llama2/runner/runner.cpp

@@ -208,7 +208,8 @@ Result<torch::executor::Tensor> Runner::run_model_step(
 Error Runner::generate(
     const std::string& prompt,
     int32_t seq_len,
-    std::function<void(const std::string&)> callback) {
+    std::function<void(const std::string&)> on_token_generated_callback,
+    std::function<void(const TimeStampsAndStats&)> on_stats_callback) {
   // Prepare the inputs.
   // Use ones-initialized inputs.
   ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null");
@@ -364,8 +365,8 @@ Error Runner::generate(
     util::safe_printf(piece);
     fflush(stdout);
 
-    if (callback) {
-      callback(piece);
+    if (on_token_generated_callback) {
+      on_token_generated_callback(piece);
     }
 
     if (shouldStop_) {
@@ -386,18 +387,21 @@ Error Runner::generate(
     ET_LOG(Info, "Sequence length (%i tokens) reached!", seq_len);
   }
 
-  timers_.printReport(num_prompt_tokens, pos - num_prompt_tokens);
+  timers_.num_prompt_tokens = num_prompt_tokens;
+  timers_.num_generated_tokens = pos - num_prompt_tokens;
+  timers_.printReport();
+  if (on_stats_callback) {
+    on_stats_callback(timers_);
+  }
 
   delete[] prompt_tokens;
   return Error::Ok;
 }
 
-void Runner::TimeStamps::printReport(
-    const int64_t& num_prompt_tokens,
-    const int64_t& num_generated_tokens) {
+void Runner::TimeStampsAndStats::printReport() {
   printf(
       "PyTorchObserver %s\n",
-      toJsonString(num_prompt_tokens, num_generated_tokens).c_str());
+      toJsonString().c_str());
 
   ET_LOG(
       Info,
@@ -449,9 +453,7 @@ void Runner::TimeStamps::printReport(
       (double)aggregate_sampling_time_ms / SCALING_FACTOR_UNITS_PER_SECOND);
 }
 
-const std::string Runner::TimeStamps::toJsonString(
-    const int64_t& num_prompt_tokens,
-    const int64_t& num_generated_tokens) {
+const std::string Runner::TimeStampsAndStats::toJsonString() {
   std::stringstream ss;
   ss << "{\"prompt_tokens\":" << num_prompt_tokens << ","
      << "\"generated_tokens\":" << num_generated_tokens << ","

examples/models/llama2/runner/runner.h

@@ -31,15 +31,7 @@ class Runner {
       const std::string& tokenizer_path,
       const float temperature = 0.8f);
 
-  bool is_loaded() const;
-  Error load();
-  Error generate(
-      const std::string& prompt,
-      int32_t seq_len = 128,
-      std::function<void(const std::string&)> callback = {});
-  void stop();
-
-  struct TimeStamps {
+  struct TimeStampsAndStats {
     // Scaling factor for timestamps - in this case, we use ms.
     const long SCALING_FACTOR_UNITS_PER_SECOND = 1000;
     // Time stamps for the different stages of the execution
@@ -59,15 +51,23 @@ class Runner {
     long inference_end_ms;
     // Keep a running total of the time spent in sampling.
     long aggregate_sampling_time_ms;
+    // Token count from prompt
+    int64_t num_prompt_tokens;
+    // Token count from generated (total - prompt)
+    int64_t num_generated_tokens;
 
-    void printReport(
-        const int64_t& num_prompt_tokens,
-        const int64_t& num_generated_tokens);
-    const std::string toJsonString(
-        const int64_t& num_prompt_tokens,
-        const int64_t& num_generated_tokens);
+    void printReport();
+    const std::string toJsonString();
   };
-  TimeStamps timers_;
+
+  bool is_loaded() const;
+  Error load();
+  Error generate(
+      const std::string& prompt,
+      int32_t seq_len = 128,
+      std::function<void(const std::string&)> on_token_generated_callback = {},
+      std::function<void(const TimeStampsAndStats&)> on_stats_callback = {});
+  void stop();
 
  private:
   // metadata
@@ -98,6 +98,7 @@ class Runner {
   std::unique_ptr<Tokenizer> tokenizer_;
   std::unique_ptr<Sampler> sampler_;
   bool shouldStop_{false};
+  TimeStampsAndStats timers_;
 };
 
 } // namespace torch::executor