Update README.md

README.md

@@ -27,6 +27,38 @@ model = RoFormerModel.from_pretrained("IDEA-CCNL/Zhouwenwang-110M")
 
 
 ```
+
+### Generate task
+You can use Zhouwenwang-110M to continue writing
+
+```python
+from model.roformer.modeling_roformer import RoFormerModel
+from transformers import AutoTokenizer
+import torch
+import numpy as np
+
+sentence = '清华大学位于'
+max_length = 32
+model_pretrained_weight_path = '/home/'  # 预训练模型权重路径
+
+tokenizer = AutoTokenizer.from_pretrained(model_pretrained_weight_path)
+model = RoFormerModel.from_pretrained(model_pretrained_weight_path)
+
+for i in range(max_length):
+    encode = torch.tensor(
+        [[tokenizer.cls_token_id]+tokenizer.encode(sentence, add_special_tokens=False)]).long()
+    logits = model(encode)[0]
+    logits = torch.nn.functional.linear(
+        logits, model.embeddings.word_embeddings.weight)
+    logits = torch.nn.functional.softmax(
+        logits, dim=-1).cpu().detach().numpy()[0]
+    sentence = sentence + \
+        tokenizer.decode(int(np.random.choice(logits.shape[1], p=logits[-1])))
+    if sentence[-1] == '。':
+        break
+print(sentence)
+```
+
 ## Scores on downstream chinese tasks (without any data augmentation)
 |     Model| afqmc    |  tnews  | iflytek    |  ocnli  |  cmnli  | wsc  | csl  |
 | :--------:    | :-----:  | :----:  | :-----:   | :----: | :----: | :----: | :----: |