How to retrieve specific matches from Trie search results in Python?

Question

I have implemented a Trie data structure in Python to search for words or phrases in long texts. The Trie contains a dataset of words/phrases with associated categories. My goal is to find specific matches in the text based on the Trie search. However, the current implementation returns more matches than expected. Here is my code:

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end_of_word = False
        self.category = None


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word, category):
        current = self.root
        for char in str(word):
            if char not in current.children:
                current.children[char] = TrieNode()
            current = current.children[char]
        current.is_end_of_word = True
        current.category = category

    def search(self, text):
        matches = []
        for i in range(len(text)):
            current = self.root
            j = i
            while j < len(text) and text[j] in current.children:
                current = current.children[text[j]]
                if current.is_end_of_word:
                    matches.append((text[i : j + 1], current.category))
                j += 1
        return matches


def build_trie(trie: Trie = None) -> Trie:
    """Builds a trie data structure.

    Args:
        trie (Trie): Trie data structure.

    Returns:
        Trie: Trie data structure.

    Raises:
        FileNotFoundError: If the file is not found.
        ValueError: If no valid data is found in the Excel file.
    """
    if trie is None:
        trie = Trie()
    try:
        categories = [
            "person",
            "organization",
            "location",
            "facility",
            "product",
            "event",
        ]
        file_names = [category + ".xlsx" for category in categories]
        entities = merge_excel_files(file_names)
        for _, row in entities.iterrows():
            trie.insert(row["word"], row["category"])
        return trie
    except Exception as e:
        LOGGER.info(msg=f"Error occurred while building the trie: {e}")
        raise


def search_entities(text: str, id: int, trie: Trie) -> dict:
    """Searches for entities in the persian text and returns the matches in JSON format for NER.

    Args:
        text (str): Text to be searched.
        trie (Trie): Trie data structure containing the entities.

    Returns:
        dict: Dictionary containing the matches:

    Examples:
        >>> search_entities('وزرارت ورزش و جوانان را دوست دارم.', trie)
            [
                text: 'وزرارت ورزش و جوانان را دوست دارم.',
                id: 5,
                pre
                {
                    id: 1,
                    "start": 0,
                    "end": 15,
                    "word": "وزارت ورزش و جوانان",
                    "category": "organization"
                },
                {
                    // ...
                }
            ]

    """
    matches = trie.search(text)
    matches = [
        {
            "text": text,
            "id": id,
            "pre": [
                {
                    "id": i + 1,
                    "start": text.find(match[0]),
                    "end": text.find(match[0]) + len(match[0]),
                    "word": match[0],
                    "category": match[1]
                }
                for i, match in enumerate(matches)
            ],
        }
    ]
    return matches


trie = build_trie()
text = "آلبرت انیشتین در بزرگراه آزادگان به همراه پرویز معین به سوی پردیس سینمایی راگا در حال حرکت هستند."
print(search_entities(text, 2500, trie))

the output is:

[
    {
        'text': 'آلبرت انیشتین در بزرگراه آزادگان به همراه پرویز معین به سوی پردیس سینمایی راگا در حال حرکت هستند.',
        'id': 2500,
        'pre': [
            {'id': 1, 'start': 0, 'end': 13, 'word': 'آلبرت انیشتین', 'category': 'person'
            },
            {'id': 2, 'start': 1, 'end': 2, 'word': 'ل', 'category': 'product'
            },
            {'id': 3, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            },
            {'id': 4, 'start': 9, 'end': 10, 'word': 'ش', 'category': 'location'
            },
            {'id': 5, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            },
            {'id': 6, 'start': 17, 'end': 32, 'word': 'بزرگراه آزادگان', 'category': 'facility'
            },
            {'id': 7, 'start': 18, 'end': 19, 'word': 'ز', 'category': 'organization'
            },
            {'id': 8, 'start': 20, 'end': 21, 'word': 'گ', 'category': 'product'
            },
            {'id': 9, 'start': 21, 'end': 24, 'word': 'راه', 'category': 'product'
            },
            {'id': 10, 'start': 25, 'end': 32, 'word': 'آزادگان', 'category': 'facility'
            },
            {'id': 11, 'start': 18, 'end': 19, 'word': 'ز', 'category': 'organization'
            },
            {'id': 12, 'start': 20, 'end': 21, 'word': 'گ', 'category': 'product'
            },
            {'id': 13, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            },
            {'id': 14, 'start': 36, 'end': 41, 'word': 'همراه', 'category': 'product'
            },
            {'id': 15, 'start': 21, 'end': 24, 'word': 'راه', 'category': 'product'
            },
            {'id': 16, 'start': 42, 'end': 45, 'word': 'پرو', 'category': 'location'
            },
            {'id': 17, 'start': 42, 'end': 52, 'word': 'پرویز معین', 'category': 'person'
            },
            {'id': 18, 'start': 43, 'end': 46, 'word': 'روی', 'category': 'organization'
            },
            {'id': 19, 'start': 18, 'end': 19, 'word': 'ز', 'category': 'organization'
            },
            {'id': 20, 'start': 48, 'end': 52, 'word': 'معین', 'category': 'person'
            },
            {'id': 21, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            },
            {'id': 22, 'start': 60, 'end': 65, 'word': 'پردیس', 'category': 'location'
            },
            {'id': 23, 'start': 60, 'end': 78, 'word': 'پردیس سینمایی راگا', 'category': 'facility'
            },
            {'id': 24, 'start': 62, 'end': 64, 'word': 'دی', 'category': 'location'
            },
            {'id': 25, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            },
            {'id': 26, 'start': 69, 'end': 71, 'word': 'ما', 'category': 'product'
            },
            {'id': 27, 'start': 69, 'end': 72, 'word': 'مای', 'category': 'organization'
            },
            {'id': 28, 'start': 20, 'end': 21, 'word': 'گ', 'category': 'product'
            },
            {'id': 29, 'start': 83, 'end': 85, 'word': 'ال', 'category': 'product'
            },
            {'id': 30, 'start': 1, 'end': 2, 'word': 'ل', 'category': 'product'
            },
            {'id': 31, 'start': 86, 'end': 88, 'word': 'حر', 'category': 'location'
            },
            {'id': 32, 'start': 7, 'end': 8, 'word': 'ن', 'category': 'product'
            }
        ]
    }
]

but I expected:

[
    {
        'text': 'آلبرت انیشتین در بزرگراه آزادگان به همراه پرویز معین به سوی پردیس سینمایی راگا در حال حرکت هستند.',
        'id': 2500,
        'pre': [
            {'id': 1, 'start': 0, 'end': 13, 'word': 'آلبرت انیشتین', 'category': 'person'
            },
            {'id': 2, 'start': 17, 'end': 32, 'word': 'بزرگراه آزادگان', 'category': 'facility'
            },
            {'id': 3, 'start': 42, 'end': 52, 'word': 'پرویز معین', 'category': 'person'
            },
            {'id': 4, 'start': 60, 'end': 78, 'word': 'پردیس سینمایی راگا', 'category': 'facility'
            }
        ]
    }
]

I would appreciate any suggestions or improvements on how to modify the search algorithm or code to achieve the desired output.