Please add a function to create a vocab.txt file from the dataset file that will be used to train the transfomer model, this function should be called and fed the dataset in the top of the code so that the "vocab.txt" file gets used in the training section, here's the code to modify: import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import os
from collections import Counter


# Check if CUDA is available and use CPU otherwise
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define a SimpleTokenizer
class SimpleTokenizer():
    def __init__(self, vocab_file):
        with open(vocab_file, "r") as f:
            words = f.readlines()
        # Add ‘UNK’ token to the vocabulary list at the end.
        words = [word.strip() for word in words] + ["UNK"]
        # Create a vocabulary mapping {token: index}
        self.vocab = {word: idx for idx, word in enumerate(words)}
        
    def encode(self, text):
        # Replace self.vocab[self.unknown_token] with self.vocab[‘UNK’] 
        # to use the index of ‘UNK’ token.
        return [self.vocab.get(token, self.vocab["UNK"]) for token in text.split()]
    
def load_dataset(file_path, tokenizer):
    with open(file_path, "r") as file:
        lines = file.readlines()
    dataset = [tokenizer.encode(line.strip()) for line in lines]
    return dataset

class SimpleTransformer(nn.Module):
    def __init__(self, num_tokens, dim_model, num_heads, num_classes, num_layers, dim_feedforward, max_seq_length):
        super(SimpleTransformer, self).__init__()
        self.token_embedding = nn.Embedding(num_tokens, dim_model)
        self.positional_encoding = nn.Parameter(torch.zeros(max_seq_length, dim_model))
        encoder_layer = nn.TransformerEncoderLayer(d_model=dim_model, nhead=num_heads, dim_feedforward=dim_feedforward)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.output_head = nn.Linear(dim_model, num_classes)

    def forward(self, src, src_mask):
        src = self.token_embedding(src) + self.positional_encoding[:src.size(1), :]
        output = self.transformer_encoder(src, src_mask)
        output = output.mean(dim=1)
        return self.output_head(output)

def generate_square_subsequent_mask(sz):
    mask = torch.triu(torch.ones(sz, sz) * float("-inf"), diagonal=1)
    return mask

# Define hyperparameters
num_tokens = 1000  # size of vocabulary
dim_model = 512  # dimension of feature vectors
num_heads = 8  # number of attention heads
num_classes = 10  # number of output classes
num_layers = 6  # number of transformer layers
dim_feedforward = 2048  # dimension of feedforward networks
max_seq_length = 100  # maximum length of input sequences

# Initialize tokenizer with your vocabulary
tokenizer = SimpleTokenizer("vocab.txt")  # Your vocab file with list of words
dataset = load_dataset("Data.txt", tokenizer)

# Convert dataset to tensor and pad to max_seq_length
src = torch.zeros((len(dataset), max_seq_length), dtype=torch.long)
for i, tokens in enumerate(dataset):
    src[i, :len(tokens)] = torch.tensor(tokens, dtype=torch.long)

# Create a transformer model instance
model = SimpleTransformer(
    num_tokens=num_tokens,
    dim_model=dim_model,
    num_heads=num_heads,
    num_classes=num_classes,
    num_layers=num_layers,
    dim_feedforward=dim_feedforward,
    max_seq_length=max_seq_length
).to(device)

# You must replace “model.transformer_encoder.src_mask.new_ones” with “generate_square_subsequent_mask”
src_mask = generate_square_subsequent_mask(max_seq_length).to(device)

# Forward pass example
output = model(src, src_mask)
print(output)

# Example training loop, replace with real data and training details
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

losses = []
accuracies = []

for epoch in range(10):
    optimizer.zero_grad()
    output = model(src, src_mask)
    target = torch.randint(num_classes, (src.size(0),)).to(device)  # Random target for demonstration
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch + 1}, Loss: {loss.item()}")

    # Calculate accuracy
    _, predicted = torch.max(output, 1)
    accuracy = (predicted == target).sum().item() / target.size(0)
    accuracies.append(accuracy)

    losses.append(loss.item())

# Don’t forget to switch model to eval mode before inference
model.eval()

# Make sure to use model(transformer_input) to get the logits for a transformer input

# Loss plot
plt.figure(figsize=(10, 5))
plt.plot(losses, label="Training loss")
plt.title("Loss over epochs")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()

# Accuracy plot
plt.figure(figsize=(10, 5))
plt.plot(accuracies, label="Training accuracy")
plt.title("Accuracy over epochs")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend()
plt.show()

def pad_and_encode(test_sentences, tokenizer, max_seq_length):
    # Convert list of test sentences to tensor and pad to max_seq_length
    src = torch.zeros((len(test_sentences), max_seq_length), dtype=torch.long)
    for i, sentence in enumerate(test_sentences):
        tokens = tokenizer.encode(sentence)
        src[i, :len(tokens)] = torch.tensor(tokens[:max_seq_length], dtype=torch.long)  # truncate if sentence is longer than max_seq_length
    return src

# Example test sentences (replace these with real test sentences)
test_sentences = ["this is a test sentence", "classification of text is the task"]
# Example true labels for test sentences (replace these with real labels)
true_labels = torch.tensor([1, 0])

# Convert test sentences to padded tensor
src_test = pad_and_encode(test_sentences, tokenizer, max_seq_length)

# Generate mask for test sentences
# Note that in a real application you would handle padding properly and create a mask that
# ignores the padding tokens during attention
test_mask = generate_square_subsequent_mask(max_seq_length).to(device)
src_test = src_test.to(device)
true_labels = true_labels.to(device)

# Perform inference
model.eval()  # Set the model to evaluation mode
with torch.no_grad():
    output = model(src_test, test_mask)
    _, predicted = torch.max(output, 1)

# Calculate test accuracy
test_accuracy = (predicted == true_labels).sum().item() / true_labels.size(0)

# Display the test results
print(f"Predicted labels: {predicted}")
print(f"True labels: {true_labels}")
print(f"Test accuracy: {test_accuracy}")

# If you want to see the probabilities or the logits
# you can apply a softmax function over the output to get the probabilities
probabilities = nn.functional.softmax(output, dim=1)
print(f"Probabilities: {probabilities}")