Update PrateritumGPT.py

PrateritumGPT.py

@@ -4,6 +4,15 @@ import torch.nn as nn
 from torch.utils.data import Dataset, DataLoader
 from torch.nn.utils.rnn import pad_sequence
 import math
+import progressbar
+
+device="cpu"
+
+def CreateBar():
+    global bar
+    bar = progressbar.ProgressBar(maxval=100, \
+    widgets=[progressbar.Bar('=', '[', ']'), ' ', progressbar.Percentage()])
+    bar.start()
 
 tokens = list("azertyuiopqsdfghjklmwxcvbnäüöß—– ")
 tokensdict = {}
@@ -31,17 +40,18 @@ class CSVDataset(Dataset):
 # Supposons que vous ayez vos données sous forme de listes
 features = []
 labels = []
+padding=len(tokens)
 
 for i in reader:
     k = []
     for j in i[2]:
-        k += [tokens.index(j)+1]
-    k += [0] * (25 - len(k))
+        k += [tokens.index(j)]
+    #k += [-1] * (25 - len(k))
     features += [torch.Tensor(k)]
     k = []
     for j in i[8]:
-        k += [tokens.index(j)+1]
-    k += [0] * (25 - len(k))
+        k += [tokens.index(j)]
+    #k += [-1] * (25 - len(k))
     labels += [torch.Tensor(k)]
 
 MyDataset = CSVDataset(features=features, labels=labels)
@@ -49,17 +59,18 @@ 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).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.custom_embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=padding).to(device)
+        self.pos_encoder = PositionalEncoding(emb_dim, dropout).to(device)
+        encoder_layer = nn.TransformerEncoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to(device)
         self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers)
-        decoder_layer = nn.TransformerDecoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to("cuda")
+        decoder_layer = nn.TransformerDecoderLayer(emb_dim, nhead, dim_feedforward, dropout, batch_first=True).to(device)
         self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers)
-        self.output_layer = nn.Linear(emb_dim, vocab_size).to("cuda")
+        self.output_layer = nn.Linear(emb_dim, vocab_size).to(device)
 
     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):
+        #print("Source:", src)
+        #print("Target:", tgt)
         src_emb = self.custom_embedding(src.long())
-        #print("Source Embedding:", src_emb.shape)
         src_emb = self.pos_encoder(src_emb)
         #print("Source Embedding:", src_emb.shape)
         tgt_emb = self.custom_embedding(tgt.long())
@@ -69,6 +80,7 @@ class TransformerModel(nn.Module):
         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[:, -1, :])
+        #print("Output:",output.shape)
         return output
 
 class PositionalEncoding(nn.Module):
@@ -89,24 +101,63 @@ class PositionalEncoding(nn.Module):
         return self.dropout(x)
     
 def collate_fn(batch):
-    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)
+    inputs = [item[0].to(device) for item in batch]
+    targets = [item[1].to(device) for item in batch]
+    inputs = pad_sequence(inputs, batch_first=True, padding_value=padding)
+    targets = pad_sequence(targets, batch_first=True, padding_value=padding)
     return inputs, targets
 
 train_loader = DataLoader(MyDataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
 
-model = TransformerModel(vocab_size=len(tokens)+1, emb_dim=32, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048)
+model = TransformerModel(vocab_size=len(tokens)+1, emb_dim=16, nhead=4, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=256)
 loss_fn = nn.CrossEntropyLoss()
 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
-epochs = 1
+epochs = 100
+
+try:
+    model.load_state_dict(torch.load("data/PrateritumGPT.pth"))
+    print("Sucessfully loaded model.")
+except:
+    pass
+
+#print(model(torch.zeros((1,25)).to(device),torch.zeros((1,25)).to(device)))
+inp=input("Which verb? ")
+src=[[]]
+tgt=[[tokens.index(inp[0])]]
+for i in inp:
+    src[0]+=[tokens.index(i)]
+str_=inp[0]
+for i in range(100):
+    out=model(torch.Tensor(src).to(device),torch.Tensor(tgt).to(device)).tolist()[0]
+    Best=0
+    Best_=tokens.index(" ")
+    for k,f in enumerate(out):
+        if f>Best:
+            Best=f
+            Best_=k
+    if Best_==len(tokens):
+        break
+    str_+=tokens[Best_]
+    tgt[0]+=[Best_]
+
+print(str_)
+
 
 for epoch in range(epochs):
     total_loss = 0.0
 
+    CreateBar()
+
+    bar.start()
+
     for batch_idx, (inputs, targets) in enumerate(train_loader):
+
+        #print("",inputs,targets)
+
+        targets.to(device)
+        inputs.to(device)
+
         for i in range(1, targets.shape[1]):
             optimizer.zero_grad()
             output = model(inputs, targets[:, :i])  # Shifted targets
@@ -117,8 +168,16 @@ for epoch in range(epochs):
 
             total_loss += loss.item()
 
-        print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {total_loss / (batch_idx + 1)}")
+            mask = targets[:, i] != len(tokens)
+            targets = targets[mask]
+            inputs = inputs[mask]
+
+        bar.update((batch_idx+1)/len(train_loader)*100)
+
+        #print(f"Epoch {epoch + 1}/{epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {total_loss / (batch_idx + 1)}")
+
+    bar.finish()
 
     print(f"Epoch {epoch + 1}/{epochs}, Loss: {total_loss / len(train_loader)}")
 
-torch.save(model, "data/PrateritumGPT.pth")
+torch.save(model.state_dict(), "data/PrateritumGPT.pth")