Notes

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

# Rest of the code remains the same as provided in the previous messages

# Instantiate your model
d_model = 512
ffn_hidden = 2048
num_heads = 8
drop_prob = 0.1
num_layers = 1
max_sequence_length = 200
vocab_size = len(combined_vocabulary) # Use the combined vocabulary

model = Transformer(d_model,
ffn_hidden,
num_heads,
drop_prob,
num_layers,
max_sequence_length,
vocab_size,
combined_to_index,
combined_to_index, # Use the same vocabulary mapping for question and answer
START_TOKEN,
END_TOKEN,
PADDING_TOKEN)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

PAD_INDEX = combined_to_index[PADDING_TOKEN]

# Define loss function and optimizer
criterion = nn.CrossEntropyLoss(ignore_index=PAD_INDEX)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# # Training loop
# num_epochs = 10
# for epoch in range(num_epochs):
# model.train()
# total_loss = 0.0

# for batch in tqdm(train_loader, desc=f"Epoch {epoch + 1}"):
# input_question = batch['input_question'].to(device)
# target_answer = batch['target_answer'].to(device)

# optimizer.zero_grad()
# output = model(input_question, target_answer)
# output_dim = output.shape[-1]
# output = output.contiguous().view(-1, output_dim)
# target_answer = target_answer.view(-1)

# loss = criterion(output, target_answer)
# loss.backward()
# optimizer.step()

# total_loss += loss.item()

# avg_loss = total_loss / len(train_loader)
# print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")


# Training loop
num_epochs = 10
for epoch in range(num_epochs):
model.train()
total_loss = 0.0

for batch in tqdm(train_loader, desc=f"Epoch {epoch + 1}"):
combined_sequence = batch['combined_sequence'].to(device) # Use 'combined_sequence' key

optimizer.zero_grad()
output = model(combined_sequence, combined_sequence) # Use the same sequence for both input and target
output_dim = output.shape[-1]
output = output.contiguous().view(-1, output_dim)
target_sequence = combined_sequence.view(-1) # Flatten the target

loss = criterion(output, target_sequence)
loss.backward()
optimizer.step()

total_loss += loss.item()

avg_loss = total_loss / len(train_loader)
print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")