Spaces:
Running on Zero
Running on Zero
| 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() | |