Update app.py

app.py

@@ -1,36 +1,7 @@
 import gradio as gr
 import io
 import numpy as np
-
-# Trie
-class TrieNode:
-    def __init__(self):
-        self.children = {}
-        self.is_end_of_token = False
-class Trie:
-    def __init__(self):
-        self.root = TrieNode()
-
-    def insert(self, token):
-        node = self.root
-        for char in token:
-            if char not in node.children:
-                node.children[char] = TrieNode()
-            node = node.children[char]
-        node.is_end_of_token = True
-
-    def search_longest_prefix(self, text, start):
-        node = self.root
-        longest_match = None
-        current_pos = start
-        
-        while current_pos < len(text) and text[current_pos] in node.children:
-            node = node.children[text[current_pos]]
-            if node.is_end_of_token:
-                longest_match = current_pos
-            current_pos += 1
-        
-        return longest_match
+from tok import Tokenizer
 
 # Vector Loader   
 def load_vectors(fname):
@@ -44,23 +15,9 @@ def load_vectors(fname):
 vectors, sorted_vector = load_vectors('wiki-news-300d-1M.vec')
 
 # Tokenizer
+tokenizer = Tokenizer(protected_words=sorted_vector)
 def tokenize(text):
-    trie = Trie()
-    for token in sorted_vector:
-        trie.insert(token)
-
-    result = []
-    start = 0
-    
-    while start < len(text):
-        longest_match = trie.search_longest_prefix(text, start)
-        if longest_match is not None:
-            result.append(text[start:longest_match+1])
-            start = longest_match + 1
-        else:
-            start += 1
-    
-    return result
+    return tokenizer.word_tokenize(text)
 
 # Interface
 def onInput(paragraph, progress = gr.Progress()):