def test_torch_encode_plus_sent_to_model(self):
import torch
from transformers import ReformerConfig, ReformerModel
# Build sequence
first_ten_tokens = list(self.big_tokenizer.get_vocab().keys())[:10]
sequence = " ".join(first_ten_tokens)
encoded_sequence = self.big_tokenizer.encode_plus(sequence,
return_tensors="pt")
batch_encoded_sequence = self.big_tokenizer.batch_encode_plus(
[sequence, sequence], return_tensors="pt")
config = ReformerConfig()
# The input gets padded during training so adjust the axial position encodings from the pretrained model value of (512, 1024)
config.axial_pos_shape = encoded_sequence["input_ids"].shape
model = ReformerModel(config)
# Reformer has config.vocab_size == tokenizer.vocab_size == len(tokenizer) - 1 = 320; len(tokenizer) is 321 (including a pad token with id 320)
assert model.get_input_embeddings(
).weight.shape[0] >= self.big_tokenizer.vocab_size
with torch.no_grad():
model(**encoded_sequence)
model(**batch_encoded_sequence)
def test_local_model_forward(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["local", "local", "local", "local"]
torch.manual_seed(0)
model = ReformerModel(ReformerConfig(**config)).to(torch_device)
model.eval()
input_ids, attn_mask = self._get_input_ids_and_mask()
hidden_states = model(input_ids=input_ids, attention_mask=attn_mask)[0]
output_slice = hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[-1.6791, 0.7171, 0.1594, 0.4063, 1.2584], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def create_and_check_reformer_model_fp16_forward(self, config, input_ids, input_mask, choice_labels):
model = ReformerModel(config=config)
model.to(torch_device)
model.half()
model.eval()
output = model(input_ids, attention_mask=input_mask)["last_hidden_state"]
self.parent.assertFalse(torch.isnan(output).any().item())
def test_lsh_model_forward(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["lsh", "lsh", "lsh", "lsh"]
config["num_buckets"] = [2, 4]
torch.manual_seed(0)
model = ReformerModel(ReformerConfig(**config)).to(torch_device)
model.eval()
input_ids, attn_mask = self._get_input_ids_and_mask()
hidden_states = model(input_ids=input_ids, attention_mask=attn_mask)[0]
output_slice = hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[-0.9896, -0.9396, -1.0831, -0.0597, 0.2456], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def __init__(self,
n_labels,
hidden_size,
dropout=0.2,
label_ignore_idx=0,
max_seq_length=128,
batch_size=32,
head_init_range=0.04,
device='cuda',
vocab_size=320):
super().__init__()
self.n_labels = n_labels
self.linear_1 = nn.Linear(hidden_size, hidden_size)
self.classification_head = nn.Linear(hidden_size, n_labels)
self.label_ignore_idx = label_ignore_idx
self.tokenizer = ReformerTokenizer.from_pretrained(
'google/reformer-crime-and-punishment')
config = ReformerConfig(
axial_pos_shape=[batch_size,
int(max_seq_length / batch_size)])
self.model = ReformerModel(config)
self.dropout = nn.Dropout(dropout)
self.device = device
# initializing classification head
self.classification_head.weight.data.normal_(mean=0.0,
std=head_init_range)
def create_and_check_reformer_model_with_attn_mask(self, config, input_ids,
input_mask, is_decoder):
# no special position embeddings
config.axial_pos_embds = False
config.is_decoder = is_decoder
if self.lsh_attn_chunk_length is not None:
# need to set chunk length equal sequence length to be certain that chunking works
config.lsh_attn_chunk_length = self.seq_length
model = ReformerModel(config=config)
model.to(torch_device)
model.eval()
# set all position encodings to zero so that postions don't matter
with torch.no_grad():
embedding = model.embeddings.position_embeddings.embedding
embedding.weight = torch.nn.Parameter(
torch.zeros(embedding.weight.shape).to(torch_device))
embedding.weight.requires_grad = False
half_seq_len = self.seq_length // 2
roll = self.chunk_length
half_input_ids = input_ids[:, :half_seq_len]
# normal padded
attn_mask = torch.cat(
[
torch.ones_like(half_input_ids),
torch.zeros_like(half_input_ids)
],
dim=-1,
)
input_ids_padded = torch.cat(
[
half_input_ids,
ids_tensor((self.batch_size, half_seq_len), self.vocab_size)
],
dim=-1,
)
# shifted padded
input_ids_roll = torch.cat(
[
half_input_ids,
ids_tensor((self.batch_size, half_seq_len), self.vocab_size)
],
dim=-1,
)
input_ids_roll = torch.roll(input_ids_roll, roll, dims=-1)
attn_mask_roll = torch.roll(attn_mask, roll, dims=-1)
output_padded = model(input_ids_padded,
attention_mask=attn_mask)[0][:, :half_seq_len]
output_padded_rolled = model(
input_ids_roll,
attention_mask=attn_mask_roll)[0][:, roll:half_seq_len + roll]
self.parent.assertTrue(
torch.allclose(output_padded, output_padded_rolled, atol=1e-3))
def create_and_check_reformer_model(self, config, input_ids, input_mask, choice_labels):
model = ReformerModel(config=config)
model.to(torch_device)
model.eval()
result = model(input_ids, attention_mask=input_mask)
result = model(input_ids)
# 2 * hidden_size because we use reversible resnet layers
self.parent.assertEqual(
result.last_hidden_state.shape, (self.batch_size, self.seq_length, 2 * self.hidden_size)
)
class AudioReformer(nn.Module):
def __init__(self, config, d_model, num_classes, dropout=0.1):
super(AudioReformer, self).__init__()
self.config = config
self.encoder = ReformerModel(self.config)
self.decoder = SimpleLinearClassifier(int(d_model * 2), num_classes,
dropout)
def forward(self, x):
x = self.encoder.forward(inputs_embeds=x)
# x = (hidden_states,) <-- tuple
# pooled means that the token is enforced to assume the whole seq meaning.
# We are interested in first hidden state
x = x[0]
pooled = x[:, 0, :]
out = self.decoder(pooled)
return out
def create_and_check_reformer_model(self, config, input_ids, input_mask, choice_labels):
model = ReformerModel(config=config)
model.to(torch_device)
model.eval()
sequence_output, _ = model(input_ids, attention_mask=input_mask)
sequence_output, _ = model(input_ids)
result = {
"sequence_output": sequence_output,
}
# 2 * hidden_size because we use reversible resnet layers
self.parent.assertListEqual(
list(result["sequence_output"].size()), [self.batch_size, self.seq_length, 2 * self.hidden_size],
)
def create_and_check_reformer_feed_forward_chunking(
self, config, input_ids, input_mask, choice_labels):
torch.manual_seed(0)
model = ReformerModel(config=config)
model.to(torch_device)
model.eval()
hidden_states_no_chunk = model(input_ids, attention_mask=input_mask)[0]
config.chunk_size_lm_head = 1
config.chunk_size_feed_forward = 1
torch.manual_seed(0)
model = ReformerModel(config=config)
model.to(torch_device)
model.eval()
hidden_states_with_chunk = model(input_ids,
attention_mask=input_mask)[0]
self.parent.assertTrue(
torch.allclose(hidden_states_no_chunk,
hidden_states_with_chunk,
atol=1e-3))
import torch
from transformers import ReformerTokenizer, ReformerModel
MODEL_MAX_LENGTH = 4608
tokenizer_config_path = "protein_reformer/spiece.model"
tokenizer = ReformerTokenizer(vocab_file=tokenizer_config_path,
do_lower_case=True,
model_max_length=MODEL_MAX_LENGTH)
model_checkpoint = 'output/checkpoint-6500/'
model = ReformerModel.from_pretrained(model_checkpoint)
sequence_file_path = "data/yeast/yeast.txt"
f = open(sequence_file_path, "r")
sequence_txt = f.readlines()
f.close()
input_sequence_list = [
tokenizer(sequence.strip(), truncation=True,
return_tensors='pt')['input_ids'].cuda()
for sequence in sequence_txt
]
model.cuda()
protein_vectors_list = [
torch.mean(model(inp)[1][-1], dim=1) for inp in input_sequence_list
]
protein_vectors = torch.cat(protein_vectors_list, dim=0)
from sklearn.manifold import TSNE
protein_vectors_tsne = TSNE(n_components=2).fit_transform(
def __init__(self, config, d_model, num_classes, dropout=0.1):
super(AudioReformer, self).__init__()
self.config = config
self.encoder = ReformerModel(self.config)
self.decoder = SimpleLinearClassifier(int(d_model * 2), num_classes,
dropout)
# from transformers import pipeline
# nlp = pipeline("sentiment-analysis")
# result = nlp("I hate you")[0]
# print(f"label: {result['label']}, with score: {round(result['score'], 4)}")
# result = nlp("I love you")[0]
# print(f"label: {result['label']}, with score: {round(result['score'], 4)}")
from transformers import ReformerTokenizer, ReformerModel
import torch
tokenizer = ReformerTokenizer.from_pretrained(
'google/reformer-crime-and-punishment')
model = ReformerModel.from_pretrained('google/reformer-crime-and-punishment',
return_dict=True)
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
outputs = model(**inputs)
last_hidden_states = outputs.last_hidden_state
print(last_hidden_states)
