fix AutoModel

tools/transformers/README-zh-Hans.md

@@ -17,9 +17,8 @@ python tools/transformers/convert2hf.py --src_folder origin_ckpt/ --tgt_folder h
 然后可以使用 `from_pretrained` 接口加载：
 
 ```python
-from modeling_internlm import InternLMForCausalLM
-
-model = InternForCausalLM.from_pretrained("hf_ckpt/")
+>>> from transformers import AutoTokenizer, AutoModel
+>>> model = AutoModel.from_pretrained("hf_ckpt/", trust_remote_code=True).cuda()
 ```
 
 

tools/transformers/README.md

@@ -16,9 +16,8 @@ python tools/transformers/convert2hf.py --src_folder origin_ckpt/ --tgt_folder h
 Then, you can load it using the `from_pretrained` interface:
 
 ```python
-from modeling_internlm import InternLMForCausalLM
-
-model = InternForCausalLM.from_pretrained("hf_ckpt/")
+>>> from transformers import AutoTokenizer, AutoModel
+>>> model = AutoModel.from_pretrained("hf_ckpt/", trust_remote_code=True).cuda()
 ```
 
 `intern_moss_example.py` demonstrates an example of how to use LoRA for fine-tuning on the `fnlp/moss-moon-002-sft` dataset.

tools/transformers/convert2hf.py

@@ -1,5 +1,6 @@
 import argparse
 import math
+import json
 import os
 import re
 import tempfile
@@ -163,6 +164,12 @@ if __name__ == "__main__":
 
     os.makedirs(target_folder, exist_ok=True)
     model.save_pretrained(target_folder, max_shard_size="20GB")
+    # TODO There should be a better way to add this.
+    with open(os.path.join(target_folder, "config.json")) as fp:
+        config_dict = json.load(fp)
+    config_dict["auto_map"]["AutoModel"] = "modeling_internlm.InternLMModel"
+    with open(os.path.join(target_folder, "config.json"), "w") as fp:
+        json.dump(config_dict, fp, indent=2)
 
     tokenizer = InternLMTokenizer(args.tokenizer)
     tokenizer.save_pretrained(target_folder)