first commit

modules/transformer_seq_encoder.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+import PytorchBoot.stereotype as stereotype
+
+
+@stereotype.module("transformer_seq_encoder")
+class TransformerSequenceEncoder(nn.Module):
+    def __init__(self, config):
+        super(TransformerSequenceEncoder, self).__init__()
+        self.config = config
+        embed_dim = config["embed_dim"]
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=config["num_heads"],
+            dim_feedforward=config["ffn_dim"],
+            batch_first=True,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=config["num_layers"]
+        )
+        self.fc = nn.Linear(embed_dim, config["output_dim"])
+
+    def encode_sequence(self, embedding_list_batch):
+        
+        lengths = []
+        
+        for embedding_list in embedding_list_batch:
+            lengths.append(len(embedding_list))
+
+        embedding_tensor = pad_sequence(embedding_list_batch, batch_first=True)  # Shape: [batch_size, max_seq_len, embed_dim]
+
+        max_len = max(lengths)
+        padding_mask = torch.tensor([([0] * length + [1] * (max_len - length)) for length in lengths], dtype=torch.bool).to(embedding_tensor.device)
+
+        transformer_output = self.transformer_encoder(embedding_tensor, src_key_padding_mask=padding_mask)
+        final_feature = transformer_output.mean(dim=1)
+        final_output = self.fc(final_feature)
+
+        return final_output
+
+
+if __name__ == "__main__":
+    config = {
+        "embed_dim": 256,
+        "num_heads": 4,  
+        "ffn_dim": 256,
+        "num_layers": 3, 
+        "output_dim": 1024, 
+    }
+
+    encoder = TransformerSequenceEncoder(config)
+    seq_len = [5, 8, 9, 4]
+    batch_size = 4
+
+    embedding_list_batch = [
+        torch.randn(seq_len[idx], config["embed_dim"]) for idx in range(batch_size)
+    ]
+    output_feature = encoder.encode_sequence(
+        embedding_list_batch
+    )
+    print("Encoded Feature:", output_feature)
+    print("Feature Shape:", output_feature.shape)