Реализация Tr ie по уровню токена слова, а не по уровню символа в Python

Question

Я пытаюсь реализовать Tr ie в Python с помощью токенов слов (в отличие от символов) и возвращать список индексов по найденным фразам. Я убираю любые знаки препинания после разбиения слов и затем пытаюсь join отступить пробелами (по существу, токенизация пробелов с пунктуацией для каждого слова). Я считаю, что мой код дает сбой в крайних случаях, когда иногда что-то не захватывается, но я не могу понять это:

import string

class Node(object):

    def __init__(self):
        self.children = {}
        self.is_leaf = False
        self.phrase_id = -1

class Trie(object):
    def __init__(self):
        self.root = Node()
        self.word_count = 0

    def get_node(self):
        return Node()

    def word_mapper(self):
        return self.word_count

    def insert(self, some_str, phrase_id):
        '''
        Insert some string into the trie if not present in trie
        If some string is prefix of trie node, mark leaf node
        '''

        search_pointer = self.root
        word_token_list = some_str.split()
        num_tokens = len(word_token_list)

        for i in range(num_tokens):

            if word_token_list[i] not in search_pointer.children:
                search_pointer.children[word_token_list[i]] = self.get_node()

            search_pointer = search_pointer.children[word_token_list[i]]
#            print("Index:", i, " Word: ", word_token_list[i])
#            print(search_pointer.children)

            if (i == num_tokens - 1):
                search_pointer.phrase_id = phrase_id

        self.word_count += 1

        # mark as leaf node for last word
        search_pointer.is_leaf = True

    def search(self, document):
        '''
        search for some string in the trie
        returns true if it exists in trie, else returns false
        '''

        found_word_list = []

        search_pointer = self.root
        doc_token_list = document.split()
        num_tokens = len(doc_token_list)

        for i in range(num_tokens):

            if doc_token_list[i] in search_pointer.children:
                search_pointer = search_pointer.children[doc_token_list[i]]
                i += 1
            else:

                # if phrase id not -1, then word found so add to list
                if (search_pointer.phrase_id != -1):
                    found_word_list.append(search_pointer.phrase_id)

                if (search_pointer == self.root):
                    i += 1
                else:
                    search_pointer = self.root

        if (search_pointer.phrase_id != -1):
            found_word_list.append(search_pointer.phrase_id)                

        return found_word_list

class DocumentProcessor(object):

    # initialize trie with word iterable as well as create a count dictionary initialized to 0
    def __init__(self, word_iterable):
        self.word_trie = Trie()

        for i in range(len(word_iterable)):
            self.word_trie.insert(word_iterable[i], i)

        self.words_list = list(word_iterable)
        self.word_dictionary = dict.fromkeys(self.words_list, 0)

    # convert word to lower case  
    def lower_caser(self, some_str):
        return some_str.lower()

    # convert word to title case
    def title_caser(self, some_str):
        return some_str.title()

    # strip punctuation from word        
    def punctuation_stripper(self, word):
        return word.strip(string.punctuation)

    # strip punctuation from whole word iterable and return as a list
    def word_iterable_punctuation_stripper(self, document):
        return [self.punctuation_stripper(self.lower_caser(word)) for word in document.split()]

    # create count of each word
    def create_count_dictionary(self, document):
        doc_word_list = self.word_iterable_punctuation_stripper(document)
        doc_rejoined = ' '.join(doc_word_list)

        for i in self.word_dictionary:
            if self.word_trie.search(doc_rejoined):
                try:
                    self.word_dictionary[self.title_caser(i)] += 1
                except KeyError:
                    continue

        return self.word_dictionary

    # return document search results
    def document_searcher(self, document):
        doc_word_list = self.word_iterable_punctuation_stripper(document)
        doc_rejoined = self.title_caser(' '.join(doc_word_list))

        # print(doc_rejoined)

        return self.word_trie.search(doc_rejoined)

    # return words in found order
    def words_lister(self, document):
        return [self.words_list[x] for x in self.document_searcher(document)]

# main function for testing

if __name__ == "__main__":

    sample_paragraph = "When we look at Number One and Number Two, we would like to look at Company One and Company Two, as well as Mary had a Little Lamb. Company One, Company Two are here."

    words_list = ["Number One", "Number Two", "Company One", "Company Two", "Mary", "Little Lamb"]

    doc_proc = DocumentProcessor(words_list)
    words_encountered = doc_proc.words_lister(sample_paragraph)

    search_res = doc_proc.document_searcher(sample_paragraph)

Вывод main выглядит следующим образом, что является крайним случаем когда несколько фраз находятся рядом друг с другом:

>>> words_encountered
['Number One',
 'Number Two',
 'Company One',
 'Company Two',
 'Mary',
 'Little Lamb',
 'Company Two']

И символ пробела в document_searcher для doc_rejoined выглядит следующим образом:

>>> print(doc_rejoined)
When We Look At Number One And Number Two We Would Like To Look At Company One And Company Two As Well As Mary Had A Little Lamb Company One Company Two Are Here

Кажется, что иногда промежуточный те не захвачены. Я пытаюсь выяснить этот крайний случай, и поскольку метод insert отлично работает в Trie, я считаю, что ошибка связана с моим движением указателя и индексацией в методе search. Однако я не могу понять, что делать.

Любая помощь будет признательна.

Спасибо!