From 59d7a1d58d2021a2891545ef296aa388d1628e37 Mon Sep 17 00:00:00 2001
From: x54-729 <17307130121@fudan.edu.cn>
Date: Mon, 10 Jul 2023 17:43:19 +0800
Subject: [PATCH] fix AutoModel

---
 tools/transformers/README-zh-Hans.md | 5 ++---
 tools/transformers/README.md         | 5 ++---
 tools/transformers/convert2hf.py     | 7 +++++++
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tools/transformers/README-zh-Hans.md b/tools/transformers/README-zh-Hans.md
index 7337424..8bbdaf5 100644
--- a/tools/transformers/README-zh-Hans.md
+++ b/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()
 ```
 
 
diff --git a/tools/transformers/README.md b/tools/transformers/README.md
index 0c4ecc8..4fe2a92 100644
--- a/tools/transformers/README.md
+++ b/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.
diff --git a/tools/transformers/convert2hf.py b/tools/transformers/convert2hf.py
index 7e86f6f..181b68c 100644
--- a/tools/transformers/convert2hf.py
+++ b/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)