Update README.md

README.md

@@ -126,12 +126,12 @@ This model was trained using the [Keras](https://keras.io/) framework. All train
 import keras
 import tensorflow as tf
 import numpy as np
+from huggingface_hub import from_pretrained_keras
 
-mlm_model = keras.models.load_model(
-    "bert_mlm.h5", custom_objects={"MaskedLanguageModel": MaskedLanguageModel}
-)
+model = from_pretrained_keras("bookbot/id-g2p-bert")
 
 MAX_LEN = 32
+MASK_TOKEN_ID = 30
 
 def inference(sequence):
     sequence = " ".join([c if c != "e" else "[mask]" for c in sequence])
@@ -140,10 +140,10 @@ def inference(sequence):
 
     tokens = tokens + pad
     input_ids = tf.convert_to_tensor(np.array([tokens]))
-    prediction = mlm_model.predict(input_ids)
+    prediction = model.predict(input_ids)
 
     # find masked idx token
-    masked_index = np.where(input_ids == mask_token_id)
+    masked_index = np.where(input_ids == MASK_TOKEN_ID)
     masked_index = masked_index[1]
 
     # get prediction at those masked index only