def main_test():
id_to_word, word_to_id = read_vocab(FLAGS.vocab_dir)
id_to_cat, cat_to_id = read_category(FLAGS.category_dir)
config = load_config(FLAGS.config_file)
model = TextRNN(config)
test(model, config, word_to_id, cat_to_id, id_to_cat)
def __init__(self):
self.config = TCNNConfig()
self.categories, self.cat_to_id = read_category()
self.words, self.word_to_id = read_vocab(vocab_dir)
self.config.vocab_size = len(self.words)
self.model = TextCNN(self.config)
self.session = tf.Session()
self.session.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess=self.session, save_path=save_path) # 读取保存的模型
def main_train():
# 如果不存在词汇表则新建
if not os.path.exists(FLAGS.vocab_dir):
build_vocab(FLAGS.train_dir, FLAGS.vocab_dir, FLAGS.vocab_size)
id_to_word, word_to_id = read_vocab(FLAGS.vocab_dir)
if not os.path.exists(FLAGS.category_dir):
build_category(FLAGS.train_dir, FLAGS.category_dir)
id_to_cat, cat_to_id = read_category(FLAGS.category_dir)
if os.path.isfile(FLAGS.config_file):
config = load_config(FLAGS.config_file)
else:
config = create_rnn_config_model(FLAGS, id_to_word)
model = TextRNN(config)
train(model, config, word_to_id, cat_to_id)
def evaluate_one(Model, config_file):
"""
加载模型进行单条样本预测
:param Model:
:param config_file:
:return:
"""
config = load_config(config_file)
id_to_word, word_to_id = read_vocab(config.vocab_dir)
id_to_cat, cat_to_id = read_category(config.category_dir)
model = Model(config)
with tf.Session() as session:
session.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(
sess=session,
save_path="checkpoints/textcnn/best_validation") # 读取保存的模型
while True:
try:
line = input("请输入测试文本: ")
if line == "exit":
exit(0)
data = [
word_to_id[x] if x in word_to_id else word_to_id["<UNK>"]
for x in line
]
pad_data = kr.preprocessing.sequence.pad_sequences(
[data],
config.seq_length,
padding="post",
truncating="post")
print(pad_data)
feed_dict = {model.input_x: pad_data, model.keep_prob: 1.0}
y_pred_cls, logits = session.run(
[model.y_pred_cls, model.logits], feed_dict=feed_dict)
print(y_pred_cls[0], tf.nn.softmax(logits).eval())
print("所属类别: {}".format(id_to_cat[y_pred_cls[0]]))
except Exception as e:
print(e)