Update inference.py

inference.py

@@ -299,7 +299,8 @@ def task_validate(seq_str, models):
 
 def task_predict(seq_str, models):
     tok, cls, ngram, elec_tok, elec_disc, glyph_map = models
-    mlm = load_bert_mlm(models[0].__class__)  # reload MLM
+    model_dir, data_dir = get_model_dir()
+    mlm = load_bert_mlm(model_dir)
     seq = parse_sequence(seq_str)
     preds = bert_predict_mask(seq, tok, mlm, top_k=5)
     print(f"\n  Input: {seq_str}")
@@ -406,4 +407,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()