Update PrateritumGPT.py

PrateritumGPT.py

@@ -49,26 +49,26 @@ MyDataset = CSVDataset(features=features, labels=labels)
 class TransformerModel(nn.Module):
     def __init__(self, vocab_size, emb_dim, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout=0.1):
         super().__init__()
-        self.custom_embedding = nn.Embedding(vocab_size, emb_dim)
-        self.pos_encoder = PositionalEncoding(emb_dim, dropout)
-        encoder_layer = nn.TransformerEncoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True)
+        self.custom_embedding = nn.Embedding(vocab_size, emb_dim).to("cuda")
+        self.pos_encoder = PositionalEncoding(emb_dim, dropout).to("cuda")
+        encoder_layer = nn.TransformerEncoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to("cuda")
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers)
-        decoder_layer = nn.TransformerDecoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True)
+        decoder_layer = nn.TransformerDecoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to("cuda")
         self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
-        self.output_layer = nn.Linear(emb_dim, vocab_size)
+        self.output_layer = nn.Linear(emb_dim, vocab_size).to("cuda")
 
     def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
         src_emb = self.custom_embedding(src.long())
-        print("Source Embedding:", src_emb.shape)
+        #print("Source Embedding:", src_emb.shape)
         src_emb = self.pos_encoder(src_emb)
-        print("Source Embedding:", src_emb.shape)
+        #print("Source Embedding:", src_emb.shape)
         tgt_emb = self.custom_embedding(tgt.long())
-        print("Target Embedding:", tgt_emb.shape)
+        #print("Target Embedding:", tgt_emb.shape)
         tgt_emb = self.pos_encoder(tgt_emb)
-        print("Target Embedding:", tgt_emb.shape)
+        #print("Target Embedding:", tgt_emb.shape)
         encoder_output = self.transformer_encoder(src_emb, src_mask, src_key_padding_mask)
         decoder_output = self.transformer_decoder(tgt_emb, encoder_output, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask)
-        output = self.output_layer(decoder_output)
+        output = self.output_layer(decoder_output[:, -1, :])
         return output
 
 class PositionalEncoding(nn.Module):
@@ -89,10 +89,10 @@ class PositionalEncoding(nn.Module):
         return self.dropout(x)
     
 def collate_fn(batch):
-    inputs = [item[0] for item in batch]
-    targets = [item[1] for item in batch]
-    inputs = pad_sequence(inputs, batch_first=True, padding_value=len(tokens) + 1)
-    targets = pad_sequence(targets, batch_first=True, padding_value=len(tokens) + 1)
+    inputs = [item[0].to("cuda") for item in batch]
+    targets = [item[1].to("cuda") for item in batch]
+    inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
+    targets = pad_sequence(targets, batch_first=True, padding_value=0)
     return inputs, targets
 
 train_loader = DataLoader(MyDataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
@@ -101,28 +101,24 @@ model = TransformerModel(vocab_size=len(tokens)+1, emb_dim=32, nhead=8, num_enco
 loss_fn = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
-epochs = 10
+epochs = 1
 
 for epoch in range(epochs):
     total_loss = 0.0
 
     for batch_idx, (inputs, targets) in enumerate(train_loader):
-
-        for i in range(1,targets.shape[1]):
-
+        for i in range(1, targets.shape[1]):
             optimizer.zero_grad()
             output = model(inputs, targets[:, :i])  # Shifted targets
-            output = output.transpose(1, 2)  # Adjust shape for loss function
-            #loss = loss_fn(output, targets[:, i].unsqueeze(1).long())  # Shifted targets
-            print(output.shape)
-            print(targets[:, i].unsqueeze(1).long().shape)
-            loss = loss_fn(output, targets[:, i].unsqueeze(1).long())
+            #print(output.shape)
+            loss = loss_fn(output, targets[:, i].long())  # Reshape targets
             loss.backward()
             optimizer.step()
 
             total_loss += loss.item()
 
-            if batch_idx % 100 == 0:
-                print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {total_loss / (batch_idx + 1)}")
+        print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {total_loss / (batch_idx + 1)}")
 
     print(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(train_loader)}")
+
+torch.save(model, "data/PrateritumGPT.pth")