Repository: piskvorky/gensim

Issue:
The keep_tokens parameter for Dictionary.filter_extremes does not work properly
The documentation of the `filter_extremes` method of `gensim.corpora.Dictionary` states for the `keep_tokens` parameter: "Iterable of tokens that **must** stay in dictionary after filtering." 

However, if filtering by the `keep_n` parameter is applied, those tokens can still get removed. 

This is due to the code around line 358 in gensim/corpora/dictionary.py where the ids are first sorted by document frequency in reverse order and then the first `keep_n` of those indices are taken. However, if the tokens listed in `keep_tokens` are rare and not within the `keep_n` most frequent ones, they get removed by this. 

The expected behaviour would be that the tokens listed in `keep_tokens` never get removed.


Discussion:

Resolving patch:
diff --git a/gensim/corpora/dictionary.py b/gensim/corpora/dictionary.py
--- a/gensim/corpora/dictionary.py
+++ b/gensim/corpora/dictionary.py
@@ -363,12 +363,13 @@ def filter_extremes(self, no_below=5, no_above=0.5, keep_n=100000, keep_tokens=N
                 v for v in itervalues(self.token2id)
                 if no_below <= self.dfs.get(v, 0) <= no_above_abs or v in keep_ids
             )
+            good_ids = sorted(good_ids, key=lambda x: self.num_docs if x in keep_ids else self.dfs.get(x), reverse=True)
         else:
             good_ids = (
                 v for v in itervalues(self.token2id)
                 if no_below <= self.dfs.get(v, 0) <= no_above_abs
             )
-        good_ids = sorted(good_ids, key=self.dfs.get, reverse=True)
+            good_ids = sorted(good_ids, key=self.dfs.get, reverse=True)
         if keep_n is not None:
             good_ids = good_ids[:keep_n]
         bad_words = [(self[idx], self.dfs.get(idx, 0)) for idx in set(self).difference(good_ids)]


Test patch:
diff --git a/gensim/test/test_corpora_dictionary.py b/gensim/test/test_corpora_dictionary.py
--- a/gensim/test/test_corpora_dictionary.py
+++ b/gensim/test/test_corpora_dictionary.py
@@ -129,6 +129,20 @@ def testFilterKeepTokens_unseenToken(self):
         expected = {'graph', 'trees', 'system', 'user'}
         self.assertEqual(set(d.token2id.keys()), expected)
 
+    def testFilterKeepTokens_keepn(self):
+        # keep_tokens should also work if the keep_n parameter is used, but only
+        # to keep a maximum of n (so if keep_n < len(keep_n) the tokens to keep are
+        # still getting removed to reduce the size to keep_n!)
+        d = Dictionary(self.texts)
+        # Note: there are four tokens with freq 3, all the others have frequence 2
+        # in self.texts. In order to make the test result deterministic, we add
+        # 2 tokens of frequency one
+        d.add_documents([['worda'], ['wordb']])
+        # this should keep the 3 tokens with freq 3 and the one we want to keep
+        d.filter_extremes(keep_n=5, no_below=0, no_above=1.0, keep_tokens=['worda'])
+        expected = {'graph', 'trees', 'system', 'user', 'worda'}
+        self.assertEqual(set(d.token2id.keys()), expected)
+
     def testFilterMostFrequent(self):
         d = Dictionary(self.texts)
         d.filter_n_most_frequent(4)