spec: save the dynamic/static ngram cache file

* fix todo on providing n_draft, save_static and save_dynamic from common/common.h

* implement the functionality by saving the cache at the common_speculative_state_ngram_cache instantiation

Author: Petros Sideris

common/common.h

@@ -308,6 +308,7 @@ struct common_params_speculative {
 
     // ngram-based speculative decoding
 
+    uint16_t ngram_n_draft  =  8; // ngram n tokens to draft
     uint16_t ngram_size_n   = 12; // ngram size for lookup
     uint16_t ngram_size_m   = 48; // mgram size for speculative tokens
     uint16_t ngram_min_hits = 1; // minimum hits at ngram/mgram lookup for mgram to be proposed
@@ -317,6 +318,9 @@ struct common_params_speculative {
     std::string lookup_cache_static;  // path of static ngram cache file for lookup decoding           // NOLINT
     std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding          // NOLINT
 
+    bool save_lookup_cache_static  = false; // whether or not we should save the static ngram cache file     // NOLINT
+    bool save_lookup_cache_dynamic = false; // whether or not we should save the dynamic ngram cache file    // NOLINT
+
     // draft-model speculative decoding
 
     struct common_params_model mparams_dft;

common/speculative.cpp

@@ -759,8 +759,6 @@ struct common_speculative_state_ngram_mod : public common_speculative_state {
 
 struct common_speculative_state_ngram_cache : public common_speculative_state {
     uint16_t n_draft;
-    bool save_dynamic;
-    bool save_static;
 
     common_ngram_cache ngram_cache_context;
     common_ngram_cache ngram_cache_dynamic;
@@ -770,19 +768,20 @@ struct common_speculative_state_ngram_cache : public common_speculative_state {
 
     common_speculative_state_ngram_cache(
             const enum common_speculative_type type,
+            uint16_t            n_draft,
             const std::string & path_static,
             const std::string & path_dynamic,
-            uint16_t            n_draft,
             bool                save_dynamic,
             bool                save_static)
         : common_speculative_state(type)
         , n_draft(n_draft)
-        , save_dynamic(save_dynamic)
-        , save_static(save_static)
     {
         if (!path_static.empty()) {
             try {
                 ngram_cache_static = common_ngram_cache_load(path_static);
+                if (save_static) {
+                    common_ngram_cache_save(ngram_cache_static, path_static);
+                }
             } catch (...) {
                 LOG_ERR("failed to open static lookup cache: %s", path_static.c_str());
                 GGML_ABORT("Couldn't read static lookup cache");
@@ -792,6 +791,9 @@ struct common_speculative_state_ngram_cache : public common_speculative_state {
         if (!path_dynamic.empty()) {
             try {
                 ngram_cache_dynamic = common_ngram_cache_load(path_dynamic);
+                if (save_dynamic) {
+                    common_ngram_cache_save(ngram_cache_dynamic, path_dynamic);
+                }
             } catch (...) {
                 LOG_ERR("failed to open dynamic lookup cache: %s", path_dynamic.c_str());
                 GGML_ABORT("Couldn't read dynamic lookup cache");
@@ -865,16 +867,15 @@ static common_ngram_map get_common_ngram_map(const common_speculative_config & c
     return common_ngram_map(size_key, size_value, key_only, min_hits);
 }
 
-static common_speculative_state_ngram_cache create_state_ngram_cache(
-        const std::string & path_static, const std::string & path_dynamic,
-        const common_speculative_config & config) {
-    uint16_t n_draft = 8; // TODO get from config?
-
-    // TODO bool param in common/common.h to set save_static/save_dynamic?
-    bool save_static = false;
-    bool save_dynamic = false;
+static common_speculative_state_ngram_cache create_state_ngram_cache(const common_speculative_config & config) {
 
-    common_speculative_state_ngram_cache state(config.type, path_static, path_dynamic, n_draft, save_static, save_dynamic);
+    common_speculative_state_ngram_cache state(
+            config.type,
+            config.params.ngram_n_draft,
+            config.params.lookup_cache_static,
+            config.params.lookup_cache_dynamic,
+            config.params.save_lookup_cache_static,
+            config.params.save_lookup_cache_dynamic);
 
     return state;
 }
@@ -1031,8 +1032,7 @@ common_speculative * common_speculative_init(
                 break;
             }
             case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE: {
-                auto state = create_state_ngram_cache(
-                        params.lookup_cache_static, params.lookup_cache_dynamic, config);
+                auto state = create_state_ngram_cache(config);
                 impls.push_back(std::make_unique<common_speculative_state_ngram_cache>(state));
                 break;
             }