Model predicted almost correct sentences at the time of training but is only predicting <START> token at the time of test

Ansh_Dadhich: <p>loss was like<br>
loss at epoch 61 is 0.01835295371711254<br>
[2, 3128, 4624, 4455, 3830, 3524, 4574, 3387, 2814, 2971, 4815, 3, 2729, 3551, 2882, 2893, 4405]<br>
&lt;&lt; START &gt;&gt; me too lets forget about germs and focus on food</p>
<p>loss at epoch 62 is 0.010977335274219513<br>
[2, 3876, 4126, 2729, 4669, 2776, 4098, 3, 2882, 4844, 3182, 2893, 4687, 3182, 4405]<br>
&lt;&lt; START &gt;&gt; then you should watch the rerun &lt;&lt; END &gt;&gt; go is end</p>
<p>loss at epoch 63 is 0.010273752734065056<br>
[2, 3324, 4844, 3987, 3, 4163, 4163, 3551, 4405, 4405, 4084, 2729, 4849, 2729, 4084, 4084, 2729]<br>
&lt;&lt; START &gt;&gt; why is that &lt;&lt; END &gt;&gt; live live where can can guy</p>
<p>loss at epoch 64 is 0.010734939947724342<br>
[2, 3987, 4844, 4319, 2776, 3466, 4844, 3083, 4139, 3, 4405, 4405, 4405, 4087, 4405, 4687]<br>
&lt;&lt; START &gt;&gt; that is true the weather is constantly changing &lt;&lt; END &gt;&gt; can</p>
<p>model was</p>
<p>import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F<br>
from torch.utils.data import DataLoader,Dataset<br>
import pandas<br>
import numpy<br>
import math<br>
import os<br>
import re<br>
import matplotlib.pyplot as mat<br>
from torch.cuda.amp import autocast,GradScaler<br>
import random</p>
<p>device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
print(device)</p>
<p>dataset = pandas.read_csv(r"Conversation.csv",usecols = [“question”,“answer”])</p>
<p>dataset_str = dataset.astype(str)<br>
length = dataset_str.map(len)<br>
dataset[“length”] = length.sum(axis = 1)</p>
<p>dataset = dataset.sort_values(by = “length”)</p>
<p>datasett = pandas.DataFrame()</p>
<p>datasett[“question”] = dataset[“question”].astype(str)<br>
datasett[“answer”] = dataset[“answer”].astype(str)</p>
<p>def remove_punctuation(sentence):<br>
sentence = re.sub(r"[^\w\s]“,”",sentence)<br>
return sentence</p>
<p>datasett[“question”] = datasett[“question”].map(remove_punctuation)<br>
datasett[“answer”] = datasett[“answer”].map(remove_punctuation)</p>
<p>vocab = [“”, “”, “”, “”] + list(set(word for sentence in datasett[“question”] for word in sentence.split())) + list(set(word for sentence in datasett[“answer”] for word in sentence.split()))</p>
<p>word_to_index = {word:index for index,word in enumerate(vocab)}<br>
index_to_word = {index:word for word,index in word_to_index.items()}</p>
<p>training_data = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<p>for i in range(0,len(dataset)):<br>
data = datasett[“question”].iloc[i]<br>
target = datasett[“answer”].iloc[i]<br>
training_data.append((data,target))</p>
<p>class custom_dataset(Dataset):<br>
def <strong>init</strong>(self,training_data,word_to_index,index_to_word):<br>
self.question,self.answer = zip(*training_data)<br>
self.word_to_index = word_to_index<br>
self.index_to_word = index_to_word</p>
<pre><code>def __len__(self):
    return len(self.question)

def __getitem__(self,index):        
    question = re.sub(r"[^\w\s]", "" ,self.question[index]).lower()
    answer = re.sub(r"[^\w\s]", "" ,self.answer[index]).lower()
    
    sentence =  torch.tensor([self.word_to_index.get("&lt;START&gt;")] + [self.word_to_index.get(word,0) for word in (question.split())] + [self.word_to_index.get("&lt;END&gt;")])
            
    target_sentence = torch.tensor([self.word_to_index.get("&lt;START&gt;")] + [self.word_to_index.get(word,0) for word in (answer.split())] + [self.word_to_index.get("&lt;END&gt;")])
                    
    return (sentence, target_sentence)
</code></pre>
<p>def collate_batch(batch):<br>
sentences,target = zip(*batch)</p>
<pre><code>padded_sentences = [x]
padded_target_sentences = []

max_sentence = len(max(sentences,key = len))
max_target = len(max(target,key = len))

max_length = max(max_sentence,max_target)

for sentence in sentences:
    padding = torch.zeros(max_length - len(sentence))
    sentence = torch.cat([sentence,padding])
    padded_sentences.append(sentence)

for target_sentence in target:
    padding = torch.zeros(max_length - len(target_sentence))
    target_sentence = torch.cat([target_sentence,padding])
    padded_target_sentences.append(target_sentence)  
    
sentences = torch.stack(padded_sentences).long()
target = torch.stack(padded_target_sentences).long()
            
return (sentences,target)
</code></pre>
<p>dataset = custom_dataset(training_data,word_to_index,index_to_word)</p>
<p>training_data = DataLoader(dataset, batch_size = 32, shuffle = True, collate_fn = collate_batch)</p>
<p>class Transformer(nn.Module):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()</p>
<pre><code>    self.embeddings = nn.Embedding(len(vocab),128)
    self.transformer = nn.Transformer(128,4,2,2,512,batch_first = True,layer_norm_eps = 0.0001)
    self.out = nn.Linear(128,len(vocab))
    
def generate_mask(self,sentence_length,):
    mask = torch.tril(torch.ones(sentence_length,sentence_length),diagonal = 0)
    return mask.masked_fill(mask == 0,-1e9)

def create_key_padding_mask(self,sequence):
     mask = (sequence != 0)
     mask = mask.float()
     return mask.masked_fill(mask == 0,-1e9)
    
def forward(self,sentence,target,src_key_padding_mask = None, target_key_padding_mask = None, no_peek_mask = None):
    sentence = self.embeddings(sentence.long())
    target = self.embeddings(target.long())
    
    output =  self.transformer(sentence,target,src_key_padding_mask = src_key_padding_mask,tgt_key_padding_mask = target_key_padding_mask, tgt_mask = no_peek_mask)
    
    output = self.out(output)
    
    return output  
</code></pre>
<p>Model = Transformer().to(device)</p>
<p>pad_index = dataset.word_to_index.get(“”)</p>
<p>loss = nn.CrossEntropyLoss(ignore_index = pad_index)</p>
<p>losses = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<p>optimizer = torch.optim.Adam(Model.parameters(), lr = 0.0001)<br>
scaler = GradScaler()</p>
<p>for epoch in range(100):<br>
Model.train()<br>
for batch,(sentence,target) in enumerate(training_data):</p>
<pre><code>    sentence = sentence.to(device)
    target = target.to(device)

    torch.autograd.set_detect_anomaly(True)
    optimizer.zero_grad()

    batch_size = sentence.size(0)
    sentence_length = sentence.size(1) 
    target_length = target.size(1)
    
    end_token_index = dataset.word_to_index.get("&lt;END&gt;")
    
    with autocast():        
         src_key_padding_mask = Model.create_key_padding_mask(sentence).to(device)
         target_key_padding_mask = Model.create_key_padding_mask(target).to(device)
         target_mask = Model.generate_mask(target_length).to(device)

         predicted = Model.forward(sentence,target,src_key_padding_mask,target_key_padding_mask,no_peek_mask = target_mask)
         
         predictedd = predicted
         
         end_mask = (target != pad_index)
         end_mask = end_mask.unsqueeze(-1).float()
                                                         
         predicted = predicted * end_mask
         
         predicted = predicted.view(-1,predicted.size(-1))
         target = target.view(-1)
 
         losss = loss(predicted,target)
 
         torch.nn.utils.clip_grad_norm_(Model.parameters(),max_norm = 0.5)
         
         scaler.scale(losss).backward()
         scaler.step(optimizer)
         scaler.update()
    
    if(losss.item() &lt; 0.003):
        break
    
losses.append(losss.detach().item())
print(f"loss at epoch {epoch} is {losss}")

Model.eval()
with torch.no_grad():
    predictedd = torch.argmax(predictedd,dim = -1).cpu().tolist()
    for sentence in predictedd:
        i = 0
        print(sentence)
        for word in sentence:
            print(dataset.index_to_word.get(word), end = " ")
            i += 1
            if(i&gt;10):
                break
        print("\n")
        break
        
if(losss.item() &lt; 0.1):
    break
</code></pre>
<p>mat.plot(range(len(losses)),losses)<br>
mat.xlabel(“epochs”)<br>
mat.ylabel(“loss”)<br>
mat.savefig(“loss.jpg”)</p>
<p>torch.save({“model_state_dict” : Model.state_dict(),<br>
“word_to_index” : dataset.word_to_index,<br>
“index_to_word” : dataset.index_to_word,<br>
},“transformer.pt”)</p>
<p>torch.cuda.empty_cache()<br>
loaded = torch.load(“transformer.pt”)</p>
<p>with torch.no_grad():<br>
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)<br>
print(device)</p>
<pre><code>predicted = "&lt;START&gt;"
last_word = ["&lt;START&gt;"]

Model.load_state_dict(loaded["model_state_dict"])
Model = Model.to(device)

for parameter in Model.parameters():
    parameter.requires_grad = False
    
word_to_index = loaded["word_to_index"]
index_to_word = loaded["index_to_word"] 

Model.eval()

torch.cuda.empty_cache()

i = 0
while (i != 10):          
       text = "hi, how are you doing?"
       text = re.sub(r"[^\w\s]","",text).lower()
       
       target = predicted
       
       target_mask = Model.generate_mask(len(target.split())).to(device)
       
       text = torch.tensor([word_to_index.get(word, word_to_index.get("&lt;UNK&gt;")) for word in text.split()]).to(device)
       output = torch.tensor([word_to_index.get(word) for word in target.split()]).to(device)
   
       text = text.unsqueeze(0).long().to(device)
       output = output.unsqueeze(0).long().to(device)
   
       predictedd = torch.argmax(Model.forward(text,target = output,no_peek_mask = target_mask),dim = -1)
                  
       if (predictedd.size(1) &gt; 1):
           predictedd = predictedd[0,-1].tolist()   
           predicted_list = index_to_word.get(predictedd)
           last_word.append(predicted_list)
           predicted = " ".join(last_word)
           print(predicted)
           torch.cuda.empty_cache()
           i += 1 
       else:
           predictedd = predictedd.squeeze().tolist()
           print(predictedd)
           last_word.append(index_to_word.get(predictedd))
           predicted = " ".join(last_word)
           print(predicted)
           i += 1
           torch.cuda.empty_cache()
        
print(predicted)

torch.cuda.empty_cache()
</code></pre>