Views
No views yet
1git clone https://github.com/qywu/memformers
2cd memformers
3pip install -e ."qywu/membart-large" (checkpooint) and "qywu/membart-base"(checkpooint).
You can directly load the checkpoint with:1import torch
2from transformers import AutoTokenizer
3from memformers.models.membart import MemBartForConditionalGeneration
4
5tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
6# load the large model in huggingface way
7membart = MemBartForConditionalGeneration.from_pretrained("qywu/membart-large")
8
9
10text1 = "Barack Obama served as the 44th President of the United States."
11text2 = "<mask> served as the 44th President of the United States."
12
13# construct the initial memory
14memory_states = membart.construct_memory(batch_size=1)
15
16# t = 0
17input_ids1 = torch.LongTensor([tokenizer.encode(text1)])
18# only run the encoder to get memory states
19encoder_outputs = membart.model.encoder(input_ids=input_ids1, memory_states=memory_states, attention_mask=None)
20memory_states = encoder_outputs.memory_states
21
22
23# t = 1
24input_ids2 = torch.LongTensor([tokenizer.encode(text2)])
25
26encoder_outputs2 = membart.model.encoder(input_ids=input_ids2, memory_states=memory_states, attention_mask=None)
27
28outputs = membart.generate(
29 encoder_outputs=encoder_outputs2,
30 decoder_start_token_id=tokenizer.bos_token_id,
31 max_length=64,
32 num_beams=1,
33 do_sample=False,
34 return_dict_in_generate=True,
35)
36
37print(tokenizer.decode(outputs.sequences[0]))
38# Barack Obama served as the 44th President of the United States.1git clone https://github.com/qywu/TorchFly
2cd TorchFly
3pip install -e .examples/finetune_dialog for details about finetuning or further pre-training MemBart on your tasks.python train.pyexamples/training_msc.1@inproceedings{DBLP:conf/ijcnlp/WuLQGGY22,
2 author = {Qingyang Wu and
3 Zhenzhong Lan and
4 Kun Qian and
5 Jing Gu and
6 Alborz Geramifard and
7 Zhou Yu},
8 title = {Memformer: {A} Memory-Augmented Transformer for Sequence Modeling},
9 booktitle = {Findings of the Association for Computational Linguistics: {AACL-IJCNLP}
10 2022, Online only, November 20-23, 2022},
11 pages = {308--318},
12 publisher = {Association for Computational Linguistics},
13 year = {2022},
14 url = {https://aclanthology.org/2022.findings-aacl.29},
15 timestamp = {Tue, 29 Nov 2022 14:53:03 +0100},
16 biburl = {https://dblp.org/rec/conf/ijcnlp/WuLQGGY22.bib},
17 bibsource = {dblp computer science bibliography, https://dblp.org}
18}1@article{DBLP:journals/corr/abs-2209-07634,
2 author = {Qingyang Wu and
3 Zhou Yu},
4 title = {Stateful Memory-Augmented Transformers for Dialogue Modeling},
5 journal = {CoRR},
6 volume = {abs/2209.07634},
7 year = {2022},
8 url = {https://doi.org/10.48550/arXiv.2209.07634},
9 doi = {10.48550/arXiv.2209.07634},
10 eprinttype = {arXiv},
11 eprint = {2209.07634},
12 timestamp = {Tue, 27 Sep 2022 16:29:43 +0200},
13 biburl = {https://dblp.org/rec/journals/corr/abs-2209-07634.bib},
14 bibsource = {dblp computer science bibliography, https://dblp.org}
15}