import json
import numpy as np
import pickle

from keras.models import load_model
from keras.preprocessing.sequence import pad_sequences

model = load_model("NER/ner_bilstm_model.keras")

with open("NER/word2idx.pkl", "rb") as f:
    word2idx = pickle.load(f)

with open("NER/tag2idx.pkl", "rb") as f:
    tag2idx = pickle.load(f)

idx2tag = {i: t for t, i in tag2idx.items()}

print(idx2tag)

maxlen = 100


def predict_sentence(sentence):
    tokens = sentence.strip().split()
    x = [word2idx.get(w.lower(), word2idx["UNK"]) for w in tokens]
    x = pad_sequences([x], maxlen=maxlen, padding="post", value=word2idx["PAD"])

    preds = model.predict(x)
    pred_labels = np.argmax(preds[0], axis=-1)

    print("Hasil prediksi NER:")
    for token, label_idx in zip(tokens, pred_labels[: len(tokens)]):
        print(f"{token}\t{idx2tag[label_idx]}")


if __name__ == "__main__":
    try:
        sentence = "saat ini indonesia adalah negara yang sangat indah"
        predict_sentence(sentence)
    except KeyboardInterrupt:
        print("\n\nSelesai.")