trie / app.py
wang-king's picture
Update app.py
13ad26f verified
import gradio as gr
from collections import defaultdict
import jieba
def get_words():
with open('./data.txt', 'r', encoding='utf-8') as f:
data = f.read()
# Use jieba for word segmentation
words = jieba.cut(data)
# Convert the result to a list or output as a string
words_list = list(words)
words_list = [word for word in words_list if word.replace(' ', '') not in ['', '\n', ',', '.', '"', 'β€œ', '”', '’', 'β€˜', '!', '-', '#', '`', '*']]
return words_list
class TrieNode:
def __init__(self):
self.children = defaultdict(TrieNode)
self.is_end_of_word = False
self.frequency = 0
class Trie:
def __init__(self):
self.root = TrieNode()
def insert(self, word):
node = self.root
for char in word:
node = node.children[char]
node.is_end_of_word = True
node.frequency += 1
def search(self, prefix):
node = self.root
for char in prefix:
if char not in node.children:
return []
node = node.children[char]
return self._find_words_from_node(node, prefix)
def _find_words_from_node(self, node, prefix):
words = []
if node.is_end_of_word:
words.append((prefix, node.frequency))
for char, next_node in node.children.items():
words.extend(self._find_words_from_node(next_node, prefix + char))
words.sort(key=lambda x: x[1], reverse=True)
return words
# Gradio interface function
def autocomplete(input_text):
if ' ' in input_text:
input_text = input_text.split(' ')[-1]
if len(input_text) == 0:
return ''
suggestions = trie.search(input_text)
# Format suggestions as a string with one word per line
formatted_suggestions = "\n".join(s[0] for s in suggestions)
# Return the formatted string
return formatted_suggestions if len(formatted_suggestions) > 0 else input_text
if __name__ == '__main__':
words = get_words()
trie = Trie()
for word in words:
trie.insert(word)
# Create Gradio interface
iface = gr.Interface(
fn=autocomplete, # Function to call
inputs=gr.Textbox(lines=1, placeholder="Type something..."), # Input textbox
outputs=gr.Textbox(lines=5, label="Suggestions"), # Output textbox
live=True, # Enable live updates
title="String Autocompletion Demo", # Interface title
description="Type a prefix to see autocomplete suggestions based on word frequency."
)
# Launch the interface
iface.launch()