doanvinhlong18/lstm-news-classifier

import pickle, json, numpy as np
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences

lstm_model = tf.keras.models.load_model('lstm_model.keras')
with open('keras_tokenizer.pkl', 'rb') as f:
    tokenizer = pickle.load(f)
with open('label_meta.json') as f:
    meta = json.load(f)

ID2LABEL   = {int(k): v for k, v in meta['id2label'].items()}
MAX_LENGTH = meta['max_length']

def predict(text):
    seq    = tokenizer.texts_to_sequences([text])
    padded = pad_sequences(seq, maxlen=MAX_LENGTH, padding='post', truncating='post')
    probs  = lstm_model.predict(padded, verbose=0).squeeze()
    idx    = int(np.argmax(probs))
    return {'label': ID2LABEL[idx], 'score': float(probs[idx])}

print(predict('AI is transforming healthcare with new diagnostic tools'))