Views
No views yet
reformer-enwik8 was pretrained on the first 90M chars of enwik8 whereas the text was chunked into batches of size 65536 chars (=2^16).
The model's weights were taken from https://console.cloud.google.com/storage/browser/trax-ml/reformer/enwik8 and converted
to Hugging Face's PyTorch ReformerLM model ReformerModelWithLMHead.1import torch
2
3# Encoding
4def encode(list_of_strings, pad_token_id=0):
5 max_length = max([len(string) for string in list_of_strings])
6
7 # create emtpy tensors
8 attention_masks = torch.zeros((len(list_of_strings), max_length), dtype=torch.long)
9 input_ids = torch.full((len(list_of_strings), max_length), pad_token_id, dtype=torch.long)
10
11 for idx, string in enumerate(list_of_strings):
12 # make sure string is in byte format
13 if not isinstance(string, bytes):
14 string = str.encode(string)
15
16 input_ids[idx, :len(string)] = torch.tensor([x + 2 for x in string])
17 attention_masks[idx, :len(string)] = 1
18
19 return input_ids, attention_masks
20
21# Decoding
22def decode(outputs_ids):
23 decoded_outputs = []
24 for output_ids in outputs_ids.tolist():
25 # transform id back to char IDs < 2 are simply transformed to ""
26 decoded_outputs.append("".join([chr(x - 2) if x > 1 else "" for x in output_ids]))
27 return decoded_outputs1from transformers import ReformerModelWithLMHead
2
3model = ReformerModelWithLMHead.from_pretrained("google/reformer-enwik8")
4encoded, attention_masks = encode(["In 1965, Brooks left IBM to found the Department of"])
5decode(model.generate(encoded, do_sample=True, max_length=150))
6
7# gives:
8# In 1965, Brooks left IBM to found the Department of Journalism in 1968. IBM had jurisdiction himself in 1980, while Brooks resolved, nevertheless thro
9ReformerModelWithLMHead is not optimized yet and is rather slow.