Update core/music_generator.py

core/music_generator.py

@@ -1,25 +1,58 @@
 # core/music_generator.py
+
 import torch
 import torch.nn as nn
 
+
 class MusicGenerator(nn.Module):
     def __init__(self, embed_dim=128, nhead=4, num_layers=4, vocab_size=128):
         super().__init__()
-        self.embedding_dim = embed_dim   # <--- ADD THIS
+
+        # expose dimensions for external use
+        self.embedding_dim = embed_dim
+        self.vocab_size = vocab_size
+
+        # layers
         self.embed = nn.Embedding(vocab_size, embed_dim)
+
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=embed_dim,
             nhead=nhead,
             batch_first=True
         )
+
         self.transformer = nn.TransformerEncoder(
             encoder_layer,
             num_layers=num_layers
         )
+
         self.fc = nn.Linear(embed_dim, vocab_size)
 
     def forward(self, x):
+        """
+        x: (B, T) integer tokens
+        """
         x = self.embed(x)
         x = self.transformer(x)
         x = self.fc(x)
-        return x
+        return x
+
+    def generate(self, seq_len=64, device="cpu"):
+        """
+        Generate a random music token sequence
+        """
+        self.eval()
+
+        # start with random token
+        tokens = torch.randint(0, self.vocab_size, (1, 1)).to(device)
+
+        for _ in range(seq_len - 1):
+            logits = self.forward(tokens)
+            next_token_logits = logits[:, -1, :]
+
+            probs = torch.softmax(next_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+
+            tokens = torch.cat([tokens, next_token], dim=1)
+
+        return tokens