Update README.md

README.md

@@ -58,19 +58,21 @@ We conducted a comprehensive evaluation of InternLM using the open-source evalua
 ### Import from Transformers
 To load the InternLM 7B Chat model using Transformers, use the following code:
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-7b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-7b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> inputs = tokenizer(["A beautiful flower"], return_tensors="pt")
->>> for k,v in inputs.items():
-        inputs[k] = v.cuda()
->>> gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
->>> output = model.generate(**inputs, **gen_kwargs)
->>> output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
->>> print(output)
-<s> A beautiful flower box made of white rose wood. It is a perfect gift for weddings, birthdays and anniversaries.
-All the roses are from our farm Roses Flanders. Therefor you know that these flowers last much longer than those in store or online!</s>
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-7b", trust_remote_code=True)
+# Set `torch_dtype=torch.float16` to load model in float16, otherwise it will be loaded as float32 and might cause OOM Error.
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-7b", torch_dtype=torch.float16, trust_remote_code=True).cuda()
+model = model.eval()
+inputs = tokenizer(["A beautiful flower"], return_tensors="pt")
+for k,v in inputs.items():
+    inputs[k] = v.cuda()
+gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
+output = model.generate(**inputs, **gen_kwargs)
+output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
+print(output)
+# <s> A beautiful flower box made of white rose wood. It is a perfect gift for weddings, birthdays and anniversaries.
+# All the roses are from our farm Roses Flanders. Therefor you know that these flowers last much longer than those in store or online!</s>
 ```
 
 ## Open Source License
@@ -109,19 +111,21 @@ InternLM ，即书生·浦语大模型，包含面向实用场景的70亿参数
 ### 通过 Transformers 加载
 通过以下的代码加载 InternLM 7B Chat 模型
 ```python
->>> from transformers import AutoTokenizer, AutoModelForCausalLM
->>> tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-7b", trust_remote_code=True)
->>> model = AutoModelForCausalLM.from_pretrained("internlm/internlm-7b", trust_remote_code=True).cuda()
->>> model = model.eval()
->>> inputs = tokenizer(["来到美丽的大自然，我们发现"], return_tensors="pt")
->>> for k,v in inputs.items():
-        inputs[k] = v.cuda()
->>> gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
->>> output = model.generate(**inputs, **gen_kwargs)
->>> output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
->>> print(output)
-来到美丽的大自然，我们发现各种各样的花千奇百怪。有的颜色鲜艳亮丽,使人感觉生机勃勃；有的是红色的花瓣儿粉嫩嫩的像少女害羞的脸庞一样让人爱不释手．有的小巧玲珑; 还有的花瓣粗大看似枯黄实则暗藏玄机！
-不同的花卉有不同的“脾气”,它们都有着属于自己的故事和人生道理.这些鲜花都是大自然中最为原始的物种,每一朵都绽放出别样的美令人陶醉、着迷!
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-7b", trust_remote_code=True)
+# `torch_dtype=torch.float16` 可以令模型以 float16 精度加载，否则 transformers 会将模型加载为 float32，有可能导致显存不足
+model = AutoModelForCausalLM.from_pretrained("internlm/internlm-7b", torch_dtype=torch.float16, trust_remote_code=True).cuda()
+model = model.eval()
+inputs = tokenizer(["来到美丽的大自然，我们发现"], return_tensors="pt")
+for k,v in inputs.items():
+    inputs[k] = v.cuda()
+gen_kwargs = {"max_length": 128, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
+output = model.generate(**inputs, **gen_kwargs)
+output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
+print(output)
+# 来到美丽的大自然，我们发现各种各样的花千奇百怪。有的颜色鲜艳亮丽,使人感觉生机勃勃；有的是红色的花瓣儿粉嫩嫩的像少女害羞的脸庞一样让人爱不释手．有的小巧玲珑; 还有的花瓣粗大看似枯黄实则暗藏玄机！
+# 不同的花卉有不同的“脾气”,它们都有着属于自己的故事和人生道理.这些鲜花都是大自然中最为原始的物种,每一朵都绽放出别样的美令人陶醉、着迷!
 ```
 
 ## 开源许可证