Update README.md

README.md

@@ -1,3 +1,131 @@
 ---
 license: apache-2.0
 ---
+
+# Memformers
+
+Memformers utilize a external dynamic memory to store history information. 
+This repo contains implementation of the pre-trained model MemBART and its training code.
+
+Check the repo [memformers](https://github.com/qywu/memformers) for details.
+
+## Install
+
+Download this repo and install it with:
+```bash
+git clone https://github.com/qywu/memformers
+cd memformers
+pip install -e .
+```
+
+## Usage
+
+
+### Inference and Generation
+
+Our implementation is based on huggingface [transformers](https://github.com/huggingface/transformers). Currently, we provide two checkpoints `"qywu/membart-large"` [(checkpooint)](https://huggingface.co/qywu/membart-large) and `"qywu/membart-base"`[(checkpooint)](https://huggingface.co/qywu/membart-base). 
+You can directly load the checkpoint with:
+
+```python
+import torch
+from transformers import AutoTokenizer
+from memformers.models.membart import MemBartForConditionalGeneration
+
+tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+# load the large model in huggingface way
+membart = MemBartForConditionalGeneration.from_pretrained("qywu/membart-large")
+
+
+text1 = "Barack Obama served as the 44th President of the United States."
+text2 = "<mask> served as the 44th President of the United States."
+
+# construct the initial memory
+memory_states = membart.construct_memory(batch_size=1)
+
+# t = 0
+input_ids1 = torch.LongTensor([tokenizer.encode(text1)])
+# only run the encoder to get memory states
+encoder_outputs = membart.model.encoder(input_ids=input_ids1, memory_states=memory_states, attention_mask=None)
+memory_states = encoder_outputs.memory_states
+
+
+# t = 1
+input_ids2 = torch.LongTensor([tokenizer.encode(text2)])
+
+encoder_outputs2 = membart.model.encoder(input_ids=input_ids2, memory_states=memory_states, attention_mask=None)
+
+outputs = membart.generate(
+    encoder_outputs=encoder_outputs2,
+    decoder_start_token_id=tokenizer.bos_token_id,
+    max_length=64,
+    num_beams=1,
+    do_sample=False,
+    return_dict_in_generate=True,
+)
+
+print(tokenizer.decode(outputs.sequences[0]))
+# Barack Obama served as the 44th President of the United States.
+```
+
+
+Note that due to [BART](https://arxiv.org/abs/1910.13461) denosing pre-training, it needs to further fine-tune the model on the downstream tasks to get better performance.
+
+### Training
+
+Training requires to install [TorchFly](https://github.com/qywu/TorchFly).
+```bash
+git clone https://github.com/qywu/TorchFly
+cd TorchFly
+pip install -e .
+```
+
+Then, you can refer to the code in `examples/finetune_dialog` for details about finetuning or further pre-training MemBart on your tasks.
+
+```python
+python train.py
+```
+
+For details, see `examples/training_msc`.
+
+## Citations
+
+Memformer: A Memory-Augmented Transformer for Sequence Modeling
+```bib
+@inproceedings{DBLP:conf/ijcnlp/WuLQGGY22,
+  author    = {Qingyang Wu and
+               Zhenzhong Lan and
+               Kun Qian and
+               Jing Gu and
+               Alborz Geramifard and
+               Zhou Yu},
+  title     = {Memformer: {A} Memory-Augmented Transformer for Sequence Modeling},
+  booktitle = {Findings of the Association for Computational Linguistics: {AACL-IJCNLP}
+               2022, Online only, November 20-23, 2022},
+  pages     = {308--318},
+  publisher = {Association for Computational Linguistics},
+  year      = {2022},
+  url       = {https://aclanthology.org/2022.findings-aacl.29},
+  timestamp = {Tue, 29 Nov 2022 14:53:03 +0100},
+  biburl    = {https://dblp.org/rec/conf/ijcnlp/WuLQGGY22.bib},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+```
+
+Stateful Memory-Augmented Transformers for Dialogue Modeling
+```bib
+@article{DBLP:journals/corr/abs-2209-07634,
+  author    = {Qingyang Wu and
+               Zhou Yu},
+  title     = {Stateful Memory-Augmented Transformers for Dialogue Modeling},
+  journal   = {CoRR},
+  volume    = {abs/2209.07634},
+  year      = {2022},
+  url       = {https://doi.org/10.48550/arXiv.2209.07634},
+  doi       = {10.48550/arXiv.2209.07634},
+  eprinttype = {arXiv},
+  eprint    = {2209.07634},
+  timestamp = {Tue, 27 Sep 2022 16:29:43 +0200},
+  biburl    = {https://dblp.org/rec/journals/corr/abs-2209-07634.bib},
+  bibsource = {dblp computer science bibliography, https://dblp.org}
+}
+```