Fix 2 of the 6 broken tests. See issue #3.

kodexlab · Jun 8, 2019 · 330bf19 · 330bf19
1 parent 907335e
commit 330bf19
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/eleve/memory.py b/eleve/memory.py
@@ -327,7 +327,9 @@ def add_sentence(self, sentence, freq=1, ngram_length=None):
             return
         if ngram_length is None:
             ngram_length = self.default_ngram_length
-        token_list = [self.sentence_start] + sentence + [self.sentence_end]
+        # We add sentence_start and sentence_end at both ends. We cast it to the type of sentence, so it
+        # works whether sentence is a str, tuple, or list.
+        token_list = type(sentence)(self.sentence_start) + sentence + type(sentence)(self.sentence_end)
         for ngram in extract_ngrams(token_list, ngram_length):
             self.fwd.add_ngram(ngram, freq)
         for ngram in extract_ngrams(token_list[::-1], ngram_length):