feat(*): change models for hg

main.py

@@ -14,11 +14,13 @@ from transformers import BertConfig, BertTokenizer, BertForSequenceClassificatio
 parser = argparse.ArgumentParser(description="Bert Chinese Text Classification")
 parser.add_argument("--mode", type=str, required=True, help="train/demo/predict")
 parser.add_argument("--data_dir", type=str, default="./data", help="training data and saved model path")
-parser.add_argument("--pretrained_bert_dir", type=str, default="./pretrained_bert", help="pretrained bert model path")
+#parser.add_argument("--pretrained_bert_dir", type=str, default="./pretrained_bert", help="pretrained bert model path")
 parser.add_argument("--seed", type=int, default=1, help="random seed for initialization")
 parser.add_argument("--input_file", type=str, default="./data/input.txt", help="input file to be predicted")
 args = parser.parse_args()
 
+args.pretrained_bert_dir = "bert-base-chinese"
+
 def set_seed(seed):
     np.random.seed(seed)
     torch.manual_seed(seed)
@@ -31,8 +33,7 @@ def main():
 
     tokenizer = BertTokenizer.from_pretrained(args.pretrained_bert_dir)
     bert_config = BertConfig.from_pretrained(args.pretrained_bert_dir, num_labels=config.num_labels)
-    model = BertForSequenceClassification.from_pretrained(
-        os.path.join(args.pretrained_bert_dir, "pytorch_model.bin"),
+    model = BertForSequenceClassification.from_pretrained(args.pretrained_bert_dir,
         config=bert_config
     )
     model.to(config.device)