Spaces:
Build error
Build error
File size: 1,964 Bytes
861c889 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
from unittest import TestCase
import torch
from transformers import AutoConfig, AutoModelWithLMHead, AutoTokenizer
import autoprompt.create_trigger as ct
def _load(model_name):
config = AutoConfig.from_pretrained('bert-base-cased')
model = AutoModelWithLMHead.from_pretrained('bert-base-cased')
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
return config, model, tokenizer
class TestGetEmbeddings(TestCase):
def test_bert(self):
model_name = 'bert-base-cased'
config, model, tokenizer = _load(model_name)
embeddings = ct.get_embeddings(model, config)
self.assertEqual(embeddings.weight.shape[0], config.vocab_size)
def test_roberta(self):
model_name = 'roberta-base'
config, model, tokenizer = _load(model_name)
embeddings = ct.get_embeddings(model, config)
self.assertEqual(embeddings.weight.shape[0], config.vocab_size)
class TestGradientStorage(TestCase):
def test_gradient_storage(self):
num_embeddings = 3
embedding_dim = 4
embeddings = torch.nn.Embedding(num_embeddings, embedding_dim)
embedding_storage = ct.GradientStorage(embeddings)
inputs = torch.tensor([0, 1, 2, 1])
outputs = embeddings(inputs)
outputs.retain_grad()
loss = outputs.sum()
loss.backward()
assert torch.equal(outputs.grad, embedding_storage.get())
def test_replace_trigger_tokens():
model_inputs = {
'input_ids': torch.tensor([
[1, 2, 3, 4],
[1, 1, 1, 0]
])
}
trigger_ids = torch.tensor([[5, 6]])
trigger_mask = torch.tensor([
[True, True, False, False],
[False, True, False, True]
])
replaced = ct.replace_trigger_tokens(model_inputs, trigger_ids, trigger_mask)
expected = torch.tensor([
[5, 6, 3, 4],
[1, 5, 1, 6]
])
assert torch.equal(expected, replaced['input_ids'])
|