Repository: bluesheeptoken/subseq (language: python)

Issue:
predict_k feature with subseq
Hi, 
Thanks for this package. I was wondering if there was an option to access the top k predictions or least get the probabilities. 
thanks !


Discussion / hints:
Hi,
Sorry for the delay. I don't know why I did not receive the notification.
Currently we cannot get the best k predictions. However, the development does not seem to be difficult.
You can give it a shot (I am here to help), or I will try to find some time to do it (this week is not feasible for me, I can try next week)

You can have a look at the CPT project which is a bit similar and the `predict_k` feature has already been asked and implemented
For reference: [predict_k CPT](https://github.com/bluesheeptoken/CPT/commit/13d6c89973b35bc863ee1d457d1e4f4c2ccfda75)
As both algorithms maintain an array of score per candidate, the implementation is probably similar

Fix patch:
diff --git a/README.md b/README.md
index 2499f33..ffea521 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,7 @@
 This project is a c++ implementation with a python wrapper of the `Succinct BWT-Based Sequence Prediction` model.
 
 Subseq is a sequence prediction model in a finite alphabet.
-It is a lossless model (does not discard information while training) and utilizes the succinct Wavelet Tree data structure and the Burrows-Wheeler Transform to compactly store and eﬃciently access training sequences for prediction.
+It is a lossless model (does not discard information while training) and utilizes the succinct Wavelet Tree data structure and the Burrows-Wheeler Transform to compactly store and efficiently access training sequences for prediction.
 
 This implementation is based on the following research paper:
 
diff --git a/subseq/cpp_sources/CSubseq.cpp b/subseq/cpp_sources/CSubseq.cpp
index e91e57e..041f8b6 100644
--- a/subseq/cpp_sources/CSubseq.cpp
+++ b/subseq/cpp_sources/CSubseq.cpp
@@ -64,7 +64,19 @@ const int CSubseq::predict_subquery(const Letters& query,
 }
 
 const int CSubseq::predict(const std::vector<int>& query) {
-    FrequencyArray frequencyArray(m_alphabet_size);
+    FrequencyArray frequency_array = compute_frequency_array(query);
+    return frequency_array.get_best_letter();
+}
+
+const std::vector<int> CSubseq::predict_k(const std::vector<int>& query,
+                                          std::size_t k) {
+    FrequencyArray frequency_array = compute_frequency_array(query);
+    return frequency_array.get_k_best_letter(k);
+}
+
+const FrequencyArray CSubseq::compute_frequency_array(
+    const std::vector<int>& query) {
+    FrequencyArray frequency_array(m_alphabet_size);
     Bitset bitset(m_fm_index.get_text_size());
     int query_used_counter = 0;
     SubqueryGenerator generator =
@@ -74,9 +86,9 @@ const int CSubseq::predict(const std::vector<int>& query) {
         if (subquery.query.empty()) break;
         query_used_counter += predict_subquery(subquery.query, query.size(),
                                                subquery.number_substitutions,
-                                               bitset, frequencyArray);
+                                               bitset, frequency_array);
         if (m_threshold_query > 0 && query_used_counter >= m_threshold_query)
             break;
     }
-    return frequencyArray.get_best_letter();
+    return frequency_array;
 }
diff --git a/subseq/cpp_sources/CSubseq.hpp b/subseq/cpp_sources/CSubseq.hpp
index 5a6458c..01d1960 100644
--- a/subseq/cpp_sources/CSubseq.hpp
+++ b/subseq/cpp_sources/CSubseq.hpp
@@ -21,8 +21,12 @@ class CSubseq {
         FrequencyArray& frequencyArray);  // return number of updates
 
     const int predict(const std::vector<int>& query);
+    const std::vector<int> predict_k(const std::vector<int>& query,
+                                     std::size_t k);
 
    private:
+    const FrequencyArray compute_frequency_array(const std::vector<int>& query);
+
     FmIndex m_fm_index;
     int m_threshold_query;
     std::size_t m_alphabet_size;
diff --git a/subseq/cpp_sources/FrequencyArray.cpp b/subseq/cpp_sources/FrequencyArray.cpp
index deafa97..fe07129 100644
--- a/subseq/cpp_sources/FrequencyArray.cpp
+++ b/subseq/cpp_sources/FrequencyArray.cpp
@@ -2,6 +2,37 @@
 
 #include <algorithm>
 
+namespace {
+struct CountingIterator {
+    typedef int value_type;
+    typedef int difference_type;
+    typedef const int& reference;
+    typedef const int* pointer;
+    typedef std::input_iterator_tag iterator_category;
+
+    CountingIterator() = default;
+    explicit CountingIterator(int value) : value(value) {}
+
+    reference operator*() const { return value; }
+
+    CountingIterator& operator++() {
+        ++value;
+        return *this;
+    }
+    CountingIterator operator++(int) { return CountingIterator(value++); }
+
+    bool operator==(const CountingIterator& rhs) const {
+        return value == rhs.value;
+    }
+    bool operator!=(const CountingIterator& rhs) const {
+        return value != rhs.value;
+    }
+
+   private:
+    int value;
+};
+}  // namespace
+
 FrequencyArray::FrequencyArray(std::size_t alphabet_size)
     : m_weights(alphabet_size) {}
 
@@ -27,3 +58,14 @@ const int FrequencyArray::get_best_letter() {
     return std::distance(m_weights.begin(),
                          std::max_element(m_weights.begin(), m_weights.end()));
 }
+
+const std::vector<int> FrequencyArray::get_k_best_letter(std::size_t k) {
+    std::vector<int> best_predictions(std::min(k, m_weights.size()));
+
+    std::partial_sort_copy(
+        CountingIterator(0), CountingIterator(m_weights.size()),
+        best_predictions.begin(), best_predictions.end(),
+        [this](int x, int y) { return m_weights[x] > m_weights[y]; });
+
+    return best_predictions;
+}
diff --git a/subseq/cpp_sources/FrequencyArray.hpp b/subseq/cpp_sources/FrequencyArray.hpp
index 3757c99..f540d4c 100644
--- a/subseq/cpp_sources/FrequencyArray.hpp
+++ b/subseq/cpp_sources/FrequencyArray.hpp
@@ -13,6 +13,7 @@ class FrequencyArray {
                                int item_index, int number_of_substitutions);
 
     const int get_best_letter();
+    const std::vector<int> get_k_best_letter(std::size_t k);
 
    private:
     std::vector<float> m_weights;
diff --git a/subseq/subseq.pyx b/subseq/subseq.pyx
index ec6c4f7..ad4e7f0 100644
--- a/subseq/subseq.pyx
+++ b/subseq/subseq.pyx
@@ -8,6 +8,7 @@ cdef extern from "cpp_sources/CSubseq.hpp":
     cdef cppclass CSubseq nogil:
         CSubseq(vector[int], int, int) except +
         int predict(vector[int])
+        vector[int] predict_k(vector[int], int)
 
 cdef class Subseq:
     cdef CSubseq *thisptr
@@ -23,7 +24,7 @@ cdef class Subseq:
         for sequence in sequences:
             for symbol in sequence:
                 text.push_back(self.alphabet.add_symbol(symbol))
-            text.push_back(0)  # Seperate sequences by symbol 0
+            text.push_back(0)  # Separate sequences by symbol 0
 
         self.thisptr = new CSubseq(text, self.alphabet.length(), self.threshold_query)
 
@@ -35,3 +36,11 @@ cdef class Subseq:
         for symbol in query:
             int_query.push_back(self.alphabet.get_index(symbol))
         return self.alphabet.get_symbol(self.thisptr.predict(int_query))
+
+    def predict_k(self, query, k):
+        cdef vector[int] int_query
+        cdef vector[int] predictions
+        for symbol in query:
+           int_query.push_back(self.alphabet.get_index(symbol))
+        predictions = self.thisptr.predict_k(int_query, k)
+        return [self.alphabet.get_symbol(index) for index in predictions]


Test patch:
diff --git a/tests/test_subseq.py b/tests/test_subseq.py
index 48da4de..d16aeaf 100644
--- a/tests/test_subseq.py
+++ b/tests/test_subseq.py
@@ -13,3 +13,8 @@ def test_predict_with_unknown_symbol(self):
         model = Subseq(0)
         model.fit(["banana"])
         self.assertEqual(model.predict("wna"), "n")
+
+    def test_predict_k(self):
+        model = Subseq(0)
+        model.fit(["ab", "ac", "ac"])
+        self.assertEqual(model.predict_k("a", 2), ["c", "b"])
diff --git a/tests_cpp/CSubseqTest.cpp b/tests_cpp/CSubseqTest.cpp
index 6a731db..6f5018b 100644
--- a/tests_cpp/CSubseqTest.cpp
+++ b/tests_cpp/CSubseqTest.cpp
@@ -1,14 +1,11 @@
 #include "gtest/gtest.h"
 #include "CSubseq.hpp"
 
-class CSubseqTest : public ::testing::Test {
- protected:
+TEST(CSubseqTest, predict) {
     Letters text = {2, 1, 3, 1, 3, 1, 0};  // banana$
     std::size_t alphabet_size = 4;
     CSubseq subseq = CSubseq(text, alphabet_size, 100);
-};
 
-TEST_F(CSubseqTest, predict) {
     Letters query;
     int expected_prediction;
     int actual_prediction;
@@ -37,3 +34,21 @@ TEST_F(CSubseqTest, predict) {
         expected_prediction
     );
 }
+
+TEST(CSubseqTest, predict_k) {
+    Letters text = {1, 2, 0, 1, 3, 0, 1, 3, 0};
+    std::size_t alphabet_size = 4;
+    CSubseq subseq = CSubseq(text, alphabet_size, 100);
+
+    Letters query = {1};
+    std::vector<int> expected = {3, 2};
+    std::vector<int> actual = subseq.predict_k(query, 2);
+
+    ASSERT_EQ(expected.size(), actual.size());
+
+    for (int i = 0; i < expected.size(); i++)
+        ASSERT_EQ(
+            expected[i],
+            actual[i]
+        );
+}
diff --git a/tests_cpp/FrequencyArrayTest.cpp b/tests_cpp/FrequencyArrayTest.cpp
index 3f9ecc8..a39bce6 100644
--- a/tests_cpp/FrequencyArrayTest.cpp
+++ b/tests_cpp/FrequencyArrayTest.cpp
@@ -19,3 +19,19 @@ TEST(FrequencyArrayTest, get_best_letter) {
     );
 }
 
+TEST(FrequencyArrayTest, get_k_best_letter) {
+    FrequencyArray frequencyArray = FrequencyArray(3);
+    frequencyArray.update(1, 3, 3, 2, 0); // lowest score
+    frequencyArray.update(2, 3, 3, 0, 0); // highest score
+
+    std::vector<int> expected = {2, 1};
+    std::vector<int> actual = frequencyArray.get_k_best_letter(2);
+
+    ASSERT_EQ(expected.size(), actual.size());
+
+    for (int i = 0; i < expected.size(); i++)
+        ASSERT_EQ(
+            expected[i],
+            actual[i]
+        );
+}