sundea/text-classification

1import argparse2import os3from importlib import import_module4 5import gradio as gr6from tqdm import tqdm7import models.TextCNN8import torch9import pickle as pkl10from utils import  build_dataset11classes=['finance','realty','stocks','education','science','society','politics','sports','game','entertainment']12 13MAX_VOCAB_SIZE = 10000  # 词表长度限制14UNK, PAD = '<UNK>', '<PAD>'  # 未知字，padding符号15def build_vocab(file_path, tokenizer, max_size, min_freq):16    vocab_dic = {}17    with open(file_path, 'r', encoding='UTF-8') as f:18        for line in tqdm(f):19            lin = line.strip()20            if not lin:21                continue22            content = lin.split('\t')[0]23            for word in tokenizer(content):24                vocab_dic[word] = vocab_dic.get(word, 0) + 125        vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[:max_size]26        vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}27        vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})28    return vocab_dic29 30parser = argparse.ArgumentParser(description='Chinese Text Classification')31parser.add_argument('--word', default=False, type=bool, help='True for word, False for char')32args = parser.parse_args()33model_name='TextCNN'34dataset = 'THUCNews'  # 数据集35embedding = 'embedding_SougouNews.npz'36x = import_module('models.' + model_name)37 38config = x.Config(dataset, embedding)39device='cuda:0'40model=models.TextCNN.Model(config)41 42# vocab, train_data, dev_data, test_data = build_dataset(config, args.word)43model.load_state_dict(torch.load('THUCNews/saved_dict/TextCNN.ckpt'))44model.to(device)45model.eval()46 47 48tokenizer = lambda x: [y for y in x]  # char-level49if os.path.exists(config.vocab_path):50    vocab = pkl.load(open(config.vocab_path, 'rb'))51else:52    vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)53    pkl.dump(vocab, open(config.vocab_path, 'wb'))54print(f"Vocab size: {len(vocab)}")55 56 57# content='时评：“国学小天才”录取缘何少佳话'58content=input('输入语句:')59 60words_line = []61token = tokenizer(content)62seq_len = len(token)63pad_size=3264contents=[]65 66if pad_size:67    if len(token) < pad_size:68        token.extend([PAD] * (pad_size - len(token)))69    else:70        token = token[:pad_size]71        seq_len = pad_size72# word to id73for word in token:74    words_line.append(vocab.get(word, vocab.get(UNK)))75 76contents.append((words_line, seq_len))77print(words_line)78# input = torch.LongTensor(words_line).unsqueeze(1).to(device)  # convert words_line to LongTensor and add batch dimension79x = torch.LongTensor([_[0] for _ in contents]).to(device)80 81        # pad前的长度(超过pad_size的设为pad_size)82seq_len = torch.LongTensor([_[1] for _ in contents]).to(device)83input=(x,seq_len)84print(input)85with torch.no_grad():86    output = model(input)87    predic = torch.max(output.data, 1)[1].cpu().numpy()88print(predic)89print('类别为：{}'.format(classes[predic[0]]))90 91 92 93 94 95# with torch.no_grad():96#     output=model(input)97# print(output)98 99#100# start_time = time.time()101# test_iter = build_iterator(test_data, config)102# with torch.no_grad():103#     predict_all = np.array([], dtype=int)104#     labels_all = np.array([], dtype=int)105#     for texts, labels in test_iter:106#         # texts=texts.to(device)107#         print(texts)108#         outputs = model(texts)109#         loss = F.cross_entropy(outputs, labels)110#         labels = labels.data.cpu().numpy()111#         predic = torch.max(outputs.data, 1)[1].cpu().numpy()112#         labels_all = np.append(labels_all, labels)113#         predict_all = np.append(predict_all, predic)114#         break115#     print(labels_all)116#     print(predict_all)117#118#119 120 121