62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
import torch.nn as nn
|
|
from typing import Any
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
from pytorch_wrapper.ptn_embedding import PtnEmbedding
|
|
|
|
|
|
class TestPtnEmbedding(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
"""Set up the test instance of PtnEmbedding."""
|
|
self.model_node = PtnEmbedding()
|
|
|
|
def test_basic_embedding(self) -> None:
|
|
"""Test that PtnEmbedding outputs expected shapes and has correct weight dimensions."""
|
|
# Create embedding module
|
|
model: nn.Module = self.model_node.f(
|
|
100, # vocabulary size
|
|
8, # embedding dimension
|
|
-1 # no padding index
|
|
)[0]
|
|
|
|
# Check embedding weight shape
|
|
self.assertEqual(model.weight.size(), torch.Size([100, 8]))
|
|
|
|
# Create sample input: batch size 1, sequence length 6
|
|
x: torch.Tensor = torch.tensor([[0, 1, 2, 3, 4, 5]], dtype=torch.int64)
|
|
|
|
# Forward pass
|
|
out: torch.Tensor = model(x)
|
|
|
|
# Output shape should be (batch_size, sequence_length, embedding_dim)
|
|
self.assertEqual(out.size(), torch.Size([1, 6, 8]))
|
|
|
|
|
|
def test_padding_index(self) -> None:
|
|
"""Test that PtnEmbedding sets the correct padding_idx internally."""
|
|
padding_idx: int = 2
|
|
model: nn.Module = self.model_node.f(
|
|
100, # vocabulary size
|
|
8, # embedding dimension
|
|
padding_idx
|
|
)[0]
|
|
|
|
# Ensure the model has the padding_idx attribute set correctly
|
|
self.assertTrue(hasattr(model, "padding_idx"))
|
|
self.assertEqual(model.padding_idx, padding_idx)
|
|
|
|
# Check that the weights at padding_idx are all zeros initially
|
|
self.assertTrue(torch.all(model.weight[padding_idx] == 0))
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|