Spaces:
Running on Zero
Running on Zero
File size: 2,406 Bytes
3d9ba5b | 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 | import unittest
import torch
from nlp_core.vocabulary import Vocabulary
from models.abstractive.attention import BahdanauAttention, LuongAttention
from models.abstractive.encoder import Seq2SeqEncoder
from models.abstractive.decoder import Seq2SeqDecoder
from models.abstractive.seq2seq_model import Seq2SeqSummarizer
class TestAbstractiveModel(unittest.TestCase):
def setUp(self):
self.src_vocab = Vocabulary("src_vocab")
self.src_vocab.add_sentence(["natural", "language", "processing", "is", "great"])
self.src_vocab.build_vocab(min_freq=1)
self.trg_vocab = Vocabulary("trg_vocab")
self.trg_vocab.add_sentence(["nlp", "is", "great"])
self.trg_vocab.build_vocab(min_freq=1)
self.emb_dim = 32
self.enc_hid_dim = 64
self.dec_hid_dim = 64
def test_attention_shapes(self):
batch_size = 2
src_len = 5
hidden = torch.randn(batch_size, self.dec_hid_dim)
encoder_outputs = torch.randn(batch_size, src_len, self.enc_hid_dim * 2)
bahdanau = BahdanauAttention(self.enc_hid_dim * 2, self.dec_hid_dim)
attn_weights = bahdanau(hidden, encoder_outputs)
self.assertEqual(attn_weights.shape, (batch_size, src_len))
self.assertAlmostEqual(attn_weights.sum(dim=1)[0].item(), 1.0, places=4)
def test_seq2seq_forward_and_decoding(self):
model = Seq2SeqSummarizer(
src_vocab=self.src_vocab,
trg_vocab=self.trg_vocab,
emb_dim=self.emb_dim,
enc_hid_dim=self.enc_hid_dim,
dec_hid_dim=self.dec_hid_dim
)
batch_size = 2
src_len = 6
trg_len = 4
src = torch.randint(0, len(self.src_vocab), (batch_size, src_len))
trg = torch.randint(0, len(self.trg_vocab), (batch_size, trg_len))
outputs = model(src, trg, teacher_forcing_ratio=0.5)
self.assertEqual(outputs.shape, (batch_size, trg_len, len(self.trg_vocab)))
# Test Greedy generation
summary_tokens, attn = model.summarize_greedy(["natural", "language", "processing"], max_len=5)
self.assertIsInstance(summary_tokens, list)
# Test Beam Search generation
beam_summary = model.summarize_beam(["natural", "language", "processing"], beam_width=2, max_len=5)
self.assertIsInstance(beam_summary, list)
if __name__ == "__main__":
unittest.main()
|