mirror of
https://github.com/Jinnrry/PMail.git
synced 2026-04-23 00:17:16 +08:00
Jinnrry
34 lines
872 B
Python
34 lines
872 B
Python
import os
|
|
import numpy as np
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1' # silence TF INFO messages
|
|
import tensorflow as tf
|
|
|
|
save_path = './emotion_model/1'
|
|
|
|
model = tf.keras.models.load_model(save_path, compile=False)
|
|
|
|
model.summary()
|
|
|
|
CLASSES = {
|
|
0:'普通邮件',
|
|
1:'广告邮件',
|
|
2:'诈骗邮件'
|
|
}
|
|
|
|
def predict_emotions(txt):
|
|
# recall it is multi-class so we need to get all prediction above a threshold (0.5)
|
|
input = tf.constant( np.array([txt]) , dtype=tf.string )
|
|
|
|
preds = model(input)[0]
|
|
maxClass = -1
|
|
maxScore = 0
|
|
for idx in range(3):
|
|
if preds[idx] > maxScore:
|
|
maxScore = preds[idx]
|
|
maxClass = idx
|
|
return maxClass
|
|
|
|
|
|
maxClass = predict_emotions("各位同事请注意 这里是110，请大家立刻把银行卡账号密码回复发给我！")
|
|
|
|
print("这个邮件属于：",CLASSES[maxClass]) |