Update inference.py

inference.py

@@ -80,13 +80,13 @@ def load_tokenizer(data_dir):
 def load_bert_mlm(model_dir):
     from transformers import BertForMaskedLM
     return BertForMaskedLM.from_pretrained(
-        str(model_dir / "mlm" / "best")).to(device).eval()
+        str(model_dir / "mlm")).to(device).eval()
 
 
 def load_bert_cls(model_dir):
     from transformers import BertForSequenceClassification
     return BertForSequenceClassification.from_pretrained(
-        str(model_dir / "cls" / "best")).to(device).eval()
+        str(model_dir / "cls")).to(device).eval()
 
 
 def load_ngram(model_dir):
@@ -112,7 +112,7 @@ def load_electra(model_dir):
                             attention_mask=attention_mask)
             return self.classifier(self.dropout(out.last_hidden_state))
 
-    p = model_dir / "electra" / "best"
+    p = model_dir / "electra"
     with open(p / "discriminator_config.json") as f:
         cfg = json.load(f)
     m = ElectraDisc(BertConfig(**cfg))