Create modeling_chord_transformer.py

modeling_chord_transformer.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+
+
+class ChordTransformerConfig(PretrainedConfig):
+    model_type = "chord-transformer"
+
+    def __init__(
+        self,
+        input_dim=12,
+        d_model=128,
+        nhead=8,
+        num_layers=4,
+        dim_feedforward=256,
+        dropout=0.1,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_layers = num_layers
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+
+
+class ChordTransformerModel(PreTrainedModel):
+    config_class = ChordTransformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.input_proj = nn.Linear(config.input_dim, config.d_model)
+        self.pos_embedding = nn.Parameter(torch.randn(1, 3000, config.d_model))
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.d_model,
+            nhead=config.nhead,
+            dim_feedforward=config.dim_feedforward,
+            dropout=config.dropout,
+            batch_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers)
+        self.output_proj = nn.Linear(config.d_model, config.input_dim)
+
+        self.post_init()
+
+    def forward(self, x, mask=None):
+        batch_size, seq_len, _ = x.size()
+        x_proj = self.input_proj(x)
+        pos_emb = self.pos_embedding[:, :seq_len, :]
+        x_proj = x_proj + pos_emb
+        output = self.encoder(x_proj, src_key_padding_mask=mask)
+        return self.output_proj(output)