# coding: UTF-8
import time
import os
import numpy as np
from trie_tree_match import Trie_tree
import random
import glob


def load_vocab(path="sensitive-words/sensitive_words_v1.txt"):
    with open(path, encoding="utf-8") as f:
        vocabs = f.read().splitlines()
        return vocabs

def filter_by_trie(text, trie_model, threhold=3):
    """
    0:dirty
    1:clean
    """
    res = trie_model.find_one(text)
    print(res)
    print(len(text))
    if len(res) > threhold:
        return 0
    elif len(res) >= 1 and len(text) < 200:
        return 0
    else:
        return 1

if __name__ == '__main__':
    dirty_vocab = load_vocab("./sensitive-words/yellow_words-v2.txt")
    trie_model = Trie_tree()
    trie_model.load_vocab(dirty_vocab)
    
    text="二、粉碎在粉碎�@道工序中，不但是把整瓶打成碎片，�有韩国19禁电影高颜值"

    label_trie = filter_by_trie(text, trie_model, 3)
    print('label:',label_trie)


