update: add new interface based on new gemma.cpp

Author: namtranase

.clang-format

@@ -0,0 +1,2 @@
+Language: Cpp
+BasedOnStyle: Google

CMakeLists.txt

@@ -11,7 +11,7 @@ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 FetchContent_Declare(sentencepiece GIT_REPOSITORY https://github.com/google/sentencepiece GIT_TAG 53de76561cfc149d3c01037f0595669ad32a5e7c)
 FetchContent_MakeAvailable(sentencepiece)
 
-FetchContent_Declare(gemma GIT_REPOSITORY https://github.com/google/gemma.cpp GIT_TAG origin/main)
+FetchContent_Declare(gemma GIT_REPOSITORY https://github.com/google/gemma.cpp GIT_TAG 8fb44ed6dd123f63dca95c20c561e8ca1de511d7)
 FetchContent_MakeAvailable(gemma)
 
 FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG da250571a45826b21eebbddc1e50d0c1137dee5f)
@@ -30,3 +30,5 @@ FetchContent_GetProperties(gemma)
 FetchContent_GetProperties(sentencepiece)
 target_include_directories(pygemma PRIVATE ${gemma_SOURCE_DIR})
 target_include_directories(pygemma PRIVATE ${sentencepiece_SOURCE_DIR})
+target_compile_definitions(libgemma PRIVATE $<$<PLATFORM_ID:Windows>:_CRT_SECURE_NO_WARNINGS NOMINMAX>)
+target_compile_options(libgemma PRIVATE $<$<PLATFORM_ID:Windows>:-Wno-deprecated-declarations>)

setup.py

@@ -3,6 +3,7 @@
 import sys
 from setuptools import setup, find_packages, Extension
 from setuptools.command.build_ext import build_ext
+import platform
 
 
 class CMakeExtension(Extension):
@@ -39,7 +40,7 @@ def build_extension(self, ext):
             "--",
             "-j",
             "12",
-        ]  # Specifies the number of jobs to run simultaneously
+        ]
 
         if not os.path.exists(self.build_temp):
             os.makedirs(self.build_temp)
@@ -58,7 +59,7 @@ def build_extension(self, ext):
     version="0.1.2",
     author="Nam Tran",
     author_email="[email protected]",
-    description="A Python package with a C++ backend using gemma.",
+    description="A Python package with a C++ backend using gemma.cpp",
     long_description="""
     This package provides Python bindings to a C++ library using pybind11.
     """,

src/gemma_binding.cpp

@@ -1,38 +1,9 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
-// Command line text interface to gemma.
-
-#include <ctime>
-#include <iostream>
-#include <random>
-#include <string>
-#include <thread>  // NOLINT
-#include <vector>
-
-// copybara:import_next_line:gemma_cpp
-#include "compression/compress.h"
-// copybara:end
-// copybara:import_next_line:gemma_cpp
-#include "gemma.h"  // Gemma
-// copybara:end
-// copybara:import_next_line:gemma_cpp
-#include "util/app.h"
-// copybara:end
-// copybara:import_next_line:gemma_cpp
-#include "util/args.h"  // HasHelp
-// copybara:end
-#include "hwy/base.h"
-#include "hwy/contrib/thread_pool/thread_pool.h"
-#include "hwy/highway.h"
-#include "hwy/per_target.h"
-#include "hwy/profiler.h"
-#include "hwy/timer.h"
-
+#include "gemma_binding.h"
 namespace py = pybind11;
 
-namespace gcpp {
-
 static constexpr std::string_view kAsciiArtBanner =
     "  __ _  ___ _ __ ___  _ __ ___   __ _   ___ _ __  _ __\n"
     " / _` |/ _ \\ '_ ` _ \\| '_ ` _ \\ / _` | / __| '_ \\| '_ \\\n"
@@ -211,35 +182,51 @@ void ReplGemma(gcpp::Gemma& model, gcpp::KVCache& kv_cache,
       << "command line flag.\n";
 }
 
-void Run(LoaderArgs& loader, InferenceArgs& inference, AppArgs& app) {
+void GemmaWrapper::loadModel(const std::vector<std::string> &args) {
+  int argc = args.size() + 1; // +1 for the program name
+  std::vector<char *> argv_vec;
+  argv_vec.reserve(argc);
+  argv_vec.push_back(const_cast<char *>("pygemma"));
+  for (const auto &arg : args)
+  {
+      argv_vec.push_back(const_cast<char *>(arg.c_str()));
+  }
+
+  char **argv = argv_vec.data();
+
+  this->m_loader = gcpp::LoaderArgs(argc, argv);
+  this->m_inference = gcpp::InferenceArgs(argc, argv);
+  this->m_app = gcpp::AppArgs(argc, argv);
+
   PROFILER_ZONE("Run.misc");
 
   hwy::ThreadPool inner_pool(0);
-  hwy::ThreadPool pool(app.num_threads);
+  hwy::ThreadPool pool(this->m_app.num_threads);
   // For many-core, pinning threads to cores helps.
-  if (app.num_threads > 10) {
-    PinThreadToCore(app.num_threads - 1);  // Main thread
+  if (this->m_app.num_threads > 10) {
+    PinThreadToCore(this->m_app.num_threads - 1);  // Main thread
 
     pool.Run(0, pool.NumThreads(),
              [](uint64_t /*task*/, size_t thread) { PinThreadToCore(thread); });
   }
 
-  gcpp::Gemma model(loader.tokenizer, loader.compressed_weights,
-                    loader.ModelType(), pool);
-
-  auto kv_cache = CreateKVCache(loader.ModelType());
+  if (!this->m_model) {
+    this->m_model.reset(new gcpp::Gemma(this->m_loader.tokenizer, this->m_loader.compressed_weights, this->m_loader.ModelType(), pool));
+  }
+//   auto kvcache = CreateKVCache(loader.ModelType());
+  this->m_kvcache = CreateKVCache(this->m_loader.ModelType());
 
-  if (const char* error = inference.Validate()) {
-    ShowHelp(loader, inference, app);
+  if (const char* error = this->m_inference.Validate()) {
+    ShowHelp(this->m_loader, this->m_inference, this->m_app);
     HWY_ABORT("\nInvalid args: %s", error);
   }
 
-  if (app.verbosity >= 1) {
+  if (this->m_app.verbosity >= 1) {
     const std::string instructions =
         "*Usage*\n"
         "  Enter an instruction and press enter (%C resets conversation, "
         "%Q quits).\n" +
-        (inference.multiturn == 0
+        (this->m_inference.multiturn == 0
              ? std::string("  Since multiturn is set to 0, conversation will "
                            "automatically reset every turn.\n\n")
              : "\n") +
@@ -252,153 +239,35 @@ void Run(LoaderArgs& loader, InferenceArgs& inference, AppArgs& app) {
 
     std::cout << "\033[2J\033[1;1H"  // clear screen
               << kAsciiArtBanner << "\n\n";
-    ShowConfig(loader, inference, app);
+    ShowConfig(this->m_loader, this->m_inference, this->m_app);
     std::cout << "\n" << instructions << "\n";
   }
-
-  ReplGemma(
-      model, kv_cache, pool, inner_pool, inference, app.verbosity,
-      /*accept_token=*/[](int) { return true; }, app.eot_line);
 }
 
-// std::string decode(gcpp::Gemma &model, hwy::ThreadPool &pool,
-//                    hwy::ThreadPool &inner_pool, const InferenceArgs &args,
-//                    int verbosity, const gcpp::AcceptFunc &accept_token,
-//                    std::string &prompt_string)
-// {
-//     std::string generated_text;
-//     // Seed the random number generator
-//     std::random_device rd;
-//     std::mt19937 gen(rd());
-//     int prompt_size{};
-//     if (model.model_training == ModelTraining::GEMMA_IT)
-//         {
-//             // For instruction-tuned models: add control tokens.
-//             prompt_string = "<start_of_turn>user\n" + prompt_string +
-//                             "<end_of_turn>\n<start_of_turn>model\n";
-//         }
-//     // Encode the prompt string into tokens
-//     std::vector<int> prompt;
-//     HWY_ASSERT(model.Tokenizer()->Encode(prompt_string, &prompt).ok());
-//     // Placeholder for generated token IDs
-//     std::vector<int> generated_tokens;
-//     // Define lambda for token decoding
-//     StreamFunc stream_token = [&generated_tokens](int token, float /* probability */) -> bool {
-//         generated_tokens.push_back(token);
-//         return true; // Continue generating
-//     };
-//     // Decode tokens
-//     prompt_size = prompt.size();
-//     GenerateGemma(model, args, prompt, /*start_pos=*/0, pool, inner_pool, stream_token,
-//                   accept_token, gen, verbosity);
-//     HWY_ASSERT(model.Tokenizer()->Decode(generated_tokens, &generated_text).ok());
-//     generated_text = generated_text.substr(prompt_string.size());
-
-// return generated_text;
-// }
-
-// std::string completion(LoaderArgs &loader, InferenceArgs &inference, AppArgs &app,
-//                        std::string &prompt_string){
-//     hwy::ThreadPool inner_pool(0);
-//     hwy::ThreadPool pool(app.num_threads);
-//     if (app.num_threads > 10)
-//     {
-//         PinThreadToCore(app.num_threads - 1); // Main thread
-
-//         pool.Run(0, pool.NumThreads(),
-//                     [](uint64_t /*task*/, size_t thread)
-//                     { PinThreadToCore(thread); });
-//     }
-//     gcpp::Gemma model(loader, pool);
-//     return decode(model, pool, inner_pool, inference, app.verbosity, /*accept_token=*/[](int)
-//                 { return true; }, prompt_string);
-
-//     }
-
-} // namespace gcpp
-
-void chat_base(int argc, char **argv)
-{
-    {
-        PROFILER_ZONE("Startup.misc");
-
-        gcpp::LoaderArgs loader(argc, argv);
-        gcpp::InferenceArgs inference(argc, argv);
-        gcpp::AppArgs app(argc, argv);
-
-        if (gcpp::HasHelp(argc, argv))
-        {
-            ShowHelp(loader, inference, app);
-            // return 0;
-        }
-
-        if (const char *error = loader.Validate())
-        {
-            ShowHelp(loader, inference, app);
-            HWY_ABORT("\nInvalid args: %s", error);
-        }
-
-        gcpp::Run(loader, inference, app);
-    }
-    PROFILER_PRINT_RESULTS(); // Must call outside the zone above.
-    // return 1;
-}
-// std::string completion_base(int argc, char **argv)
-// {
-//     gcpp::LoaderArgs loader(argc, argv);
-//     gcpp::InferenceArgs inference(argc, argv);
-//     gcpp::AppArgs app(argc, argv);
-//     std::string prompt_string = argv[argc-1];
-//     return gcpp::completion(loader, inference, app, prompt_string);
-// }
-// std::string completion_base_wrapper(const std::vector<std::string> &args,std::string &prompt_string)
-// {
-//     int argc = args.size() + 2; // +1 for the program name
-//     std::vector<char *> argv_vec;
-//     argv_vec.reserve(argc);
-
-//     argv_vec.push_back(const_cast<char *>("pygemma"));
-
-//     for (const auto &arg : args)
-//     {
-//         argv_vec.push_back(const_cast<char *>(arg.c_str()));
-//     }
-//     argv_vec.push_back(const_cast<char *>(prompt_string.c_str()));
-//     char **argv = argv_vec.data();
-//     return completion_base(argc, argv);
-// }
-void show_help_wrapper()
-{
-    // Assuming ShowHelp does not critically depend on argv content
-    gcpp::LoaderArgs loader(0, nullptr);
-    gcpp::InferenceArgs inference(0, nullptr);
-    gcpp::AppArgs app(0, nullptr);
-
-    ShowHelp(loader, inference, app);
+void GemmaWrapper::showConfig() {
+    ShowConfig(this->m_loader,this->m_inference, this->m_app);
 }
 
-std::string chat_base_wrapper(const std::vector<std::string> &args)
-{
-    int argc = args.size() + 1; // +1 for the program name
-    std::vector<char *> argv_vec;
-    argv_vec.reserve(argc);
-    argv_vec.push_back(const_cast<char *>("pygemma"));
-
-    for (const auto &arg : args)
-    {
-        argv_vec.push_back(const_cast<char *>(arg.c_str()));
-    }
-
-    char **argv = argv_vec.data();
-
-    chat_base(argc, argv);
+void GemmaWrapper::showHelp() {
+    ShowHelp(this->m_loader,this->m_inference, this->m_app);
 }
 
 
-PYBIND11_MODULE(pygemma, m)
-{
-    m.doc() = "Pybind11 integration for chat_base function";
-    m.def("chat_base", &chat_base_wrapper, "A wrapper for the chat_base function accepting Python list of strings as arguments");
-    m.def("show_help", &show_help_wrapper, "A wrapper for show_help function");
-    // m.def("completion", &completion_base_wrapper, "A wrapper for inference function");
+PYBIND11_MODULE(pygemma, m) {
+    py::class_<GemmaWrapper>(m, "Gemma")
+        .def(py::init<>())
+        .def("show_config", &GemmaWrapper::showConfig)
+        .def("show_help",  &GemmaWrapper::showHelp)
+        .def("load_model", [](GemmaWrapper &self,
+                              const std::string &tokenizer,
+                              const std::string &compressed_weights,
+                              const std::string &model) {
+            std::vector<std::string> args = {
+                "--tokenizer", tokenizer,
+                "--compressed_weights", compressed_weights,
+                "--model", model
+            };
+            self.loadModel(args); // Assuming GemmaWrapper::loadModel accepts std::vector<std::string>
+        }, py::arg("tokenizer"), py::arg("compressed_weights"), py::arg("model"))
+        .def("completion", &GemmaWrapper::completionPrompt);
 }

src/gemma_binding.h

@@ -0,0 +1,46 @@
+#pragma once
+// Command line text interface to gemma.
+
+#include <ctime>
+#include <iostream>
+#include <random>
+#include <string>
+#include <thread>  // NOLINT
+#include <vector>
+
+// copybara:import_next_line:gemma_cpp
+#include "compression/compress.h"
+// copybara:end
+// copybara:import_next_line:gemma_cpp
+#include "gemma.h"  // Gemma
+// copybara:end
+// copybara:import_next_line:gemma_cpp
+#include "util/app.h"
+// copybara:end
+// copybara:import_next_line:gemma_cpp
+#include "util/args.h"  // HasHelp
+// copybara:end
+#include "hwy/base.h"
+#include "hwy/contrib/thread_pool/thread_pool.h"
+#include "hwy/highway.h"
+#include "hwy/per_target.h"
+#include "hwy/profiler.h"
+#include "hwy/timer.h"
+
+using namespace gcpp;
+
+class GemmaWrapper {
+    public:
+      // GemmaWrapper();
+      void loadModel(const std::vector<std::string> &args); // Consider exception safety
+      void showConfig();
+      void showHelp();
+      std::string completionPrompt();
+
+    private:
+        gcpp::LoaderArgs m_loader = gcpp::LoaderArgs(0, nullptr);
+        gcpp::InferenceArgs m_inference = gcpp::InferenceArgs(0, nullptr);
+        gcpp::AppArgs m_app = gcpp::AppArgs(0, nullptr);
+        std::unique_ptr<gcpp::Gemma> m_model;
+        KVCache m_kvcache;
+};