Upload ChatIPC.cpp

ChatIPC.cpp

@@ -468,20 +468,32 @@ struct PtrEq   { bool operator()(StrPtr a, StrPtr b) const noexcept { return a =
 
 using NextSet = std::vector<StrPtr>;
 
+struct NgramHash {
+    std::size_t operator()(const std::vector<StrPtr>& v) const noexcept {
+        std::size_t seed = v.size();
+        for(auto& p : v) {
+            seed ^= std::hash<StrPtr>()(p) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+        }
+        return seed;
+    }
+};
+
 struct KnowledgeBase {
     StringInterner interner;
-    std::unordered_map<StrPtr, NextSet, PtrHash, PtrEq> next;
-    std::unordered_map<std::string, StrPtr> next_key_index;
+    // Modified to use a vector of StrPtr (N-gram) as the key
+    std::unordered_map<std::vector<StrPtr>, NextSet, NgramHash> next;
     mutable std::mutex m;
 
     std::unordered_map<StrPtr, std::vector<StrPtr>, PtrHash, PtrEq> def_index;
     mutable std::mutex def_m;
     int def_depth = 0;
 
-    void add_pair_interned(StrPtr k, StrPtr v){
+    void add_ngram(const std::vector<StrPtr>& ctx, StrPtr v){
+        for (auto p : ctx) ensure_def_for_interned(p);
+        ensure_def_for_interned(v);
+
         std::lock_guard<std::mutex> lk(m);
-        next_key_index.emplace(*k, k);
-        auto &vec = next[k];
+        auto &vec = next[ctx];
         for (auto p : vec) if (p == v) return;
         vec.push_back(v);
     }
@@ -495,8 +507,7 @@ struct KnowledgeBase {
     }
 
     void ensure_def_for_interned(StrPtr wp){
-        if (wp == nullptr) return;
-        if (def_depth <= 0) return;
+        if (wp == nullptr || def_depth <= 0) return;
 
         {
             std::lock_guard<std::mutex> lk(def_m);
@@ -546,30 +557,26 @@ struct KnowledgeBase {
         }
     }
 
-    void add_pair(const std::string &k, const std::string &v){
-        StrPtr kp = interner.intern(k);
-        StrPtr vp = interner.intern(v);
-        ensure_def_for_interned(kp);
-        ensure_def_for_interned(vp);
-        add_pair_interned(kp, vp);
-    }
-
-    std::optional<NextSet> lookup_by_string(const std::string &k) const {
+    std::optional<NextSet> lookup_ngram(const std::vector<StrPtr>& ctx) const {
         std::lock_guard<std::mutex> lk(m);
-        auto kit = next_key_index.find(k);
-        if (kit == next_key_index.end()) return std::nullopt;
-        auto it = next.find(kit->second);
+        auto it = next.find(ctx);
         if (it == next.end()) return std::nullopt;
         return it->second;
     }
+};
 
-    std::optional<NextSet> lookup_by_ptr(StrPtr k) const {
-        std::lock_guard<std::mutex> lk(m);
-        auto it = next.find(k);
-        if (it == next.end()) return std::nullopt;
-        return it->second;
+static void learn_tokens_ngram(KnowledgeBase &kb, const std::vector<std::string>& tokens, int n_gram_size) {
+    std::vector<StrPtr> window;
+    for (const auto& tok : tokens) {
+        StrPtr tp = kb.interner.intern(tok);
+        for (size_t j = 0; j < window.size(); ++j) {
+            std::vector<StrPtr> ctx(window.begin() + j, window.end());
+            kb.add_ngram(ctx, tp);
+        }
+        window.push_back(tp);
+        if (window.size() > static_cast<size_t>(n_gram_size)) window.erase(window.begin());
     }
-};
+}
 
 static std::vector<StrPtr>
 intern_tokens(KnowledgeBase &kb, const std::vector<std::string> &tokens)
@@ -872,14 +879,18 @@ static std::string best_candidate_by_similarity(
 static std::vector<std::string> construct_response(KnowledgeBase &kb,
                                                    const std::vector<std::string> &prompt_toks,
                                                    size_t response_maxlen,
-                                                   double repeat_penalty)
+                                                   double repeat_penalty,
+                                                   int n_gram_size)
 {
     std::vector<std::string> resp;
     if (prompt_toks.empty() || response_maxlen == 0) return resp;
 
     auto prompt_ptrs = intern_tokens(kb, prompt_toks);
     std::vector<StrPtr> resp_ptrs;
-    std::unordered_map<std::string,int> recent_counts;
+    std::unordered_map<std::string, int> recent_counts;
+
+    // The context combines prompt and generated text
+    std::vector<StrPtr> context = prompt_ptrs;
 
     auto would_create_2_cycle = [&](const std::string &cand) -> bool {
         if (resp.size() < 3) return false;
@@ -887,40 +898,37 @@ static std::vector<std::string> construct_response(KnowledgeBase &kb,
                normalize_dictionary_key(resp.back()) == normalize_dictionary_key(resp[resp.size() - 3]);
     };
 
-    std::string last_printed;
-
     for (size_t step = 0; step < response_maxlen; ++step){
         NextSet candidates;
         bool found = false;
         std::string context_tok;
 
-        if (step == 0){
-            for (ssize_t p = static_cast<ssize_t>(prompt_toks.size()) - 1; p >= 0; --p){
-                auto opt = kb.lookup_by_string(prompt_toks[(size_t)p]);
-                if (opt){
+        // Dynamic Backoff N-gram lookup
+        int current_n = std::min(static_cast<int>(context.size()), n_gram_size);
+        while (current_n > 0) {
+            std::vector<StrPtr> search_ctx(context.end() - current_n, context.end());
+            auto opt = kb.lookup_ngram(search_ctx);
+            if (opt && !opt->empty()) {
+                candidates = *opt;
+                found = true;
+                context_tok = *search_ctx.back();
+                break;
+            }
+            current_n--;
+        }
+
+        // Fallback to checking earlier prompt tokens if sequence stalled
+        if (!found) {
+            for (ssize_t p = static_cast<ssize_t>(prompt_ptrs.size()) - 1; p >= 0; --p){
+                std::vector<StrPtr> search_ctx = { prompt_ptrs[(size_t)p] };
+                auto opt = kb.lookup_ngram(search_ctx);
+                if (opt && !opt->empty()){
                     candidates = *opt;
                     found = true;
-                    context_tok = prompt_toks[(size_t)p];
+                    context_tok = *prompt_ptrs[(size_t)p];
                     break;
                 }
             }
-        } else {
-            auto opt = kb.lookup_by_string(last_printed);
-            if (opt){
-                candidates = *opt;
-                found = true;
-                context_tok = last_printed;
-            } else {
-                for (ssize_t p = static_cast<ssize_t>(prompt_toks.size()) - 1; p >= 0; --p){
-                    auto opt2 = kb.lookup_by_string(prompt_toks[(size_t)p]);
-                    if (opt2){
-                        candidates = *opt2;
-                        found = true;
-                        context_tok = prompt_toks[(size_t)p];
-                        break;
-                    }
-                }
-            }
         }
 
         if (!found || candidates.empty()) break;
@@ -931,9 +939,10 @@ static std::vector<std::string> construct_response(KnowledgeBase &kb,
             if (recent_counts[only_key.empty() ? only : only_key] > 0) break;
 
             resp.push_back(only);
-            resp_ptrs.push_back(kb.interner.intern(only));
+            StrPtr ptr = kb.interner.intern(only);
+            resp_ptrs.push_back(ptr);
+            context.push_back(ptr);
             recent_counts[only_key.empty() ? only : only_key] += 1;
-            last_printed = only;
             std::cout << only << ' ' << std::flush;
             continue;
         }
@@ -947,37 +956,37 @@ static std::vector<std::string> construct_response(KnowledgeBase &kb,
         if (would_create_2_cycle(chosen)) break;
 
         resp.push_back(chosen);
-        resp_ptrs.push_back(kb.interner.intern(chosen));
+        StrPtr chosen_ptr = kb.interner.intern(chosen);
+        resp_ptrs.push_back(chosen_ptr);
+        context.push_back(chosen_ptr);
 
         std::string chosen_key = normalize_dictionary_key(chosen);
         recent_counts[chosen_key.empty() ? chosen : chosen_key] += 1;
 
-        last_printed = chosen;
         std::cout << chosen << ' ' << std::flush;
     }
 
     return resp;
 }
 
-static void learn_from_file(KnowledgeBase &kb, const std::string &fname){
+static void learn_from_file(KnowledgeBase &kb, const std::string &fname, int n_gram_size){
     std::ifstream ifs(fname);
     if (!ifs) return;
+    std::vector<std::string> tokens;
     std::string tok;
-    std::string prev;
-    bool have_prev = false;
-    while (ifs >> tok){
-        if (have_prev) kb.add_pair(prev, tok);
-        prev = tok; have_prev = true;
-    }
+    while (ifs >> tok) tokens.push_back(tok);
+    learn_tokens_ngram(kb, tokens, n_gram_size);
 }
 
-static void learn_files_parallel(KnowledgeBase &kb, const std::vector<std::string> &files){
+static void learn_files_parallel(KnowledgeBase &kb, const std::vector<std::string> &files, int n_gram_size){
     #pragma omp parallel for schedule(dynamic)
-    for (ptrdiff_t i=0;i<static_cast<ptrdiff_t>(files.size());++i) learn_from_file(kb, files[(size_t)i]);
+    for (ptrdiff_t i=0;i<static_cast<ptrdiff_t>(files.size());++i){
+        learn_from_file(kb, files[(size_t)i], n_gram_size);
+    }
 }
 
 static constexpr std::uint64_t KB_MAGIC   = 0x434850434B535641ULL;
-static constexpr std::uint64_t KB_VERSION = 1ULL;
+static constexpr std::uint64_t KB_VERSION = 2ULL;
 
 static void write_u64(std::ostream &os, std::uint64_t v){
     os.write(reinterpret_cast<const char*>(&v), sizeof(v));
@@ -1001,7 +1010,7 @@ static void write_string(std::ostream &os, const std::string &s){
 
 static std::string read_string(std::istream &is){
     std::uint64_t n = read_u64(is);
-    if (n > (1ULL << 30)) throw std::runtime_error("corrupt save file: string too large");
+    if (n > (1ULL << 30)) throw std::runtime_error("save file is corrupted: string too large");
 
     std::string s;
     s.resize(static_cast<size_t>(n));
@@ -1041,7 +1050,10 @@ static void save_kb_binary(const KnowledgeBase &kb, const std::string &fname){
 
         write_u64(ofs, static_cast<std::uint64_t>(kb.next.size()));
         for (const auto &pr : kb.next){
-            write_u64(ofs, id.at(*pr.first));
+            write_u64(ofs, static_cast<std::uint64_t>(pr.first.size()));
+            for (StrPtr ctx_tok : pr.first){
+                write_u64(ofs, id.at(*ctx_tok));
+            }
             write_u64(ofs, static_cast<std::uint64_t>(pr.second.size()));
             for (StrPtr nxt : pr.second){
                 write_u64(ofs, id.at(*nxt));
@@ -1081,7 +1093,7 @@ static void load_kb_binary(KnowledgeBase &kb, const std::string &fname, int cli_
     const std::uint64_t file_def_depth = read_u64(ifs);
 
     const std::uint64_t N = read_u64(ifs);
-    if (N > (1ULL << 26)) throw std::runtime_error("corrupt save file: pool too large");
+    if (N > (1ULL << 26)) throw std::runtime_error("save file is corrupted: pool too large");
 
     std::vector<std::string> strings;
     strings.reserve(static_cast<size_t>(N));
@@ -1101,43 +1113,47 @@ static void load_kb_binary(KnowledgeBase &kb, const std::string &fname, int cli_
 
     // Rebuild next
     const std::uint64_t E = read_u64(ifs);
-    if (E > (1ULL << 26)) throw std::runtime_error("corrupt save file: graph too large");
+    if (E > (1ULL << 26)) throw std::runtime_error("save file is corrupted: graph too large");
 
     {
         std::lock_guard<std::mutex> lk(kb.m);
         kb.next.clear();
-        kb.next_key_index.clear();
         kb.next.reserve(static_cast<size_t>(E));
-        kb.next_key_index.reserve(static_cast<size_t>(E));
     }
 
     for (std::uint64_t i = 0; i < E; ++i){
-        const std::uint64_t key_idx = read_u64(ifs);
-        const std::uint64_t M = read_u64(ifs);
+        const std::uint64_t ctx_size = read_u64(ifs);
+        if (ctx_size > (1ULL << 16)) throw std::runtime_error("save file is corrupted: n-gram context too large");
+
+        std::vector<StrPtr> ctx;
+        ctx.reserve(static_cast<size_t>(ctx_size));
+        for(std::uint64_t c = 0; c < ctx_size; ++c) {
+            const std::uint64_t c_idx = read_u64(ifs);
+            if (c_idx >= ptrs.size()) throw std::runtime_error("save file is corrupted: bad context key");
+            ctx.push_back(ptrs[(size_t)c_idx]);
+        }
 
-        if (key_idx >= ptrs.size()) throw std::runtime_error("corrupt save file: bad graph key");
-        if (M > (1ULL << 26)) throw std::runtime_error("corrupt save file: graph degree too large");
+        const std::uint64_t M = read_u64(ifs);
+        if (M > (1ULL << 26)) throw std::runtime_error("save file is corrupted: graph degree too large");
 
-        StrPtr key_ptr = ptrs[(size_t)key_idx];
         NextSet vec;
         vec.reserve(static_cast<size_t>(M));
 
         for (std::uint64_t j = 0; j < M; ++j){
             const std::uint64_t v_idx = read_u64(ifs);
-            if (v_idx >= ptrs.size()) throw std::runtime_error("corrupt save file: bad graph value");
+            if (v_idx >= ptrs.size()) throw std::runtime_error("save file is corrupted: bad graph value");
             vec.push_back(ptrs[(size_t)v_idx]);
         }
 
         {
             std::lock_guard<std::mutex> lk(kb.m);
-            kb.next.emplace(key_ptr, std::move(vec));
-            kb.next_key_index.emplace(*key_ptr, key_ptr);
+            kb.next.emplace(std::move(ctx), std::move(vec));
         }
     }
 
     // Rebuild def_index from file
     const std::uint64_t K = read_u64(ifs);
-    if (K > (1ULL << 26)) throw std::runtime_error("corrupt save file: def_index too large");
+    if (K > (1ULL << 26)) throw std::runtime_error("save file is corrupted: def_index too large");
 
     {
         std::lock_guard<std::mutex> lk(kb.def_m);
@@ -1150,15 +1166,15 @@ static void load_kb_binary(KnowledgeBase &kb, const std::string &fname, int cli_
         const std::uint64_t key_idx = read_u64(ifs);
         const std::uint64_t M = read_u64(ifs);
 
-        if (key_idx >= ptrs.size()) throw std::runtime_error("corrupt save file: bad def key");
-        if (M > (1ULL << 26)) throw std::runtime_error("corrupt save file: def list too large");
+        if (key_idx >= ptrs.size()) throw std::runtime_error("save file is corrupted: bad def key");
+        if (M > (1ULL << 26)) throw std::runtime_error("save file is corrupted: def list too large");
 
         std::vector<StrPtr> toks;
         toks.reserve(static_cast<size_t>(M));
 
         for (std::uint64_t j = 0; j < M; ++j){
             const std::uint64_t v_idx = read_u64(ifs);
-            if (v_idx >= ptrs.size()) throw std::runtime_error("corrupt save file: bad def value");
+            if (v_idx >= ptrs.size()) throw std::runtime_error("save file is corrupted: bad def value");
             toks.push_back(ptrs[(size_t)v_idx]);
         }
 
@@ -1184,9 +1200,16 @@ static void load_kb_binary(KnowledgeBase &kb, const std::string &fname, int cli_
         {
             std::lock_guard<std::mutex> lk(kb.m);
             for (const auto &pr : kb.next){
-                if (seen.insert(pr.first).second) targets.push_back(pr.first);
+                for (StrPtr ctx_element : pr.first) {
+                    if (seen.insert(ctx_element).second) {
+                        targets.push_back(ctx_element);
+                    }
+                }
+
                 for (StrPtr v : pr.second){
-                    if (seen.insert(v).second) targets.push_back(v);
+                    if (seen.insert(v).second) {
+                        targets.push_back(v);
+                    }
                 }
             }
         }
@@ -1201,20 +1224,22 @@ static void load_kb_binary(KnowledgeBase &kb, const std::string &fname, int cli_
 static void print_usage(const char *p){
     std::cout << "CLI options: [--response-max-length N] [--save FILE] [--load-kb FILE] [--dictionary-depth D] [--learn f1 f2 ...] [--repeat-penalty P] [--help]\n";
     std::cout << "  --response-max-length N     Maximum number of tokens in a response.\n";
-    std::cout << "  --save FILE                 Save the knowledge-base to a binary file.\n";
+    std::cout << "  --save-kb FILE              Save the knowledge-base to a binary file.\n";
     std::cout << "  --load-kb FILE              Load a previously saved knowledge-base from a binary file.\n";
     std::cout << "  --dictionary-depth D        Depth of dictionary-definition expansion used during learning.\n";
+    std::cout << "  --n-gram N                  Size of the n-gram where N is the size.\n";
     std::cout << "  --learn f1 f2 ...           Learn from one or more text files to update the knowledge-base.\n";
     std::cout << "  --repeat-penalty P          Penalize repeated tokens when constructing response (higher values reduce repetition).\n";
-    std::cout << "  --help                      Show Command-Line Interface (CLI) options for using " << p << "\n";
+    std::cout << "  --help                      Show all the Command-Line Interface (CLI) options for using " << p << "\n";
 }
 
 int main(int argc, char **argv){
-    size_t response_maxlen = 100;
+    size_t response_maxlen = 500;
     std::string savefile;
     std::string load_txt;
     std::string load_kb;
-    int dict_depth = 2;
+    int dict_depth = 3;
+    int n_gram_size = 3;
     double repeat_penalty = 0.7; // default λ
     std::vector<std::string> learn_files;
 
@@ -1222,18 +1247,13 @@ int main(int argc, char **argv){
         std::string a = argv[i];
         if (a=="--help"){ print_usage(argv[0]); return 0; }
         if (a=="--response-max-length" && i+1<argc){ response_maxlen = std::stoul(argv[++i]); continue; }
-        if (a=="--save" && i+1<argc){ savefile = argv[++i]; continue; }
+        if (a=="--save-kb" && i+1<argc){ savefile = argv[++i]; continue; }
         if (a=="--load-kb" && i+1<argc){ load_kb = argv[++i]; continue; }
         if (a=="--dictionary-depth" && i+1<argc){ dict_depth = std::stoi(argv[++i]); continue; }
+        if (a=="--n-gram" && i+1<argc){ n_gram_size = std::max(1, std::stoi(argv[++i])); continue; }
         if (a=="--repeat-penalty" && i+1<argc){ repeat_penalty = std::stod(argv[++i]); continue; }
         if (a=="--learn"){
-            ++i;
-            for (; i<argc; ++i){
-                if (!argv[i]) break;
-                std::string s = argv[i];
-                if (!s.empty() && s[0]=='-'){ --i; break; }
-                learn_files.push_back(s);
-            }
+            while(i+1<argc && argv[i+1][0] != '-') learn_files.push_back(argv[++i]);
             continue;
         }
         learn_files.push_back(a);
@@ -1253,7 +1273,7 @@ int main(int argc, char **argv){
 
     if (!learn_files.empty()){
         std::cerr << "Learning from file/s (" << learn_files.size() << ") using threads=" << omp_get_max_threads() << "\n";
-        learn_files_parallel(kb, learn_files);
+        learn_files_parallel(kb, learn_files, n_gram_size);
     }
 
     std::string line;
@@ -1261,10 +1281,10 @@ int main(int argc, char **argv){
     while (std::cout << "> " , std::getline(std::cin, line)){
         if (line.empty()){ std::cout << "\n"; continue; }
         auto prompt_toks = tokenize_whitespace(line);
-        for (size_t i=1;i<prompt_toks.size();++i) kb.add_pair(prompt_toks[i-1], prompt_toks[i]);
-        auto resp = construct_response(kb, prompt_toks, response_maxlen, repeat_penalty);
+        learn_tokens_ngram(kb, prompt_toks, n_gram_size);
+        auto resp = construct_response(kb, prompt_toks, response_maxlen, repeat_penalty, n_gram_size);
         std::cout << "\n";
-        if (!resp.empty()){for (size_t i=1;i<resp.size();++i) kb.add_pair(resp[i-1], resp[i]);}
+        if (!resp.empty()){learn_tokens_ngram(kb, resp, n_gram_size);}
         if (!savefile.empty()){
             try { std::cerr << "Saving knowledge base: " << savefile << "\n";
             save_kb_binary(kb, savefile); std::cerr << "Saved knowledge base: " << savefile << "\n"; }