Upload 6 files

load_data.py

@@ -0,0 +1,71 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import DataLoader, Dataset
+from collections import Counter
+from itertools import chain
+
+
+# # Hyperparameters
+# d_model = 512  # Dimension of the embeddings and the token representations
+seq_length = 10  # Length of the input and output sequences
+# vocab_size = 1000  # Size of the vocabulary
+batch_size = 32  # Batch size for training
+# num_heads = 8  # Number of heads in multi-head attention
+# dim_feedforward = 2048  # Dimension of feedforward network in encoder and decoder
+
+
+# Assuming `transformer_model` and hyperparameters are defined as before
+
+class TextDataset(Dataset):
+    def __init__(self, text, vocab=None, seq_length=seq_length):
+        # Tokenization - simple split by whitespace
+        self.tokens = text.split()
+
+        # If a vocabulary is provided, use it, otherwise create a new one
+        if vocab:
+            self.vocab = vocab
+        else:
+            # Build vocabulary from the unique tokens with added <pad> and <eos> tokens
+            self.vocab = {'<pad>': 0, '<eos>': 1}
+            token_counts = Counter(self.tokens)
+            for token, _ in token_counts.items():
+                self.vocab[token] = len(self.vocab)
+        
+        # Inverse mapping from indices to tokens
+        self.index2token = {index: token for token, index in self.vocab.items()}
+
+        # Convert tokens to indices
+        self.indexed_tokens = [self.vocab[token] for token in self.tokens]
+        
+        # Sequence length
+        self.seq_length = seq_length
+    
+    def __len__(self):
+        # Number of tokens divided by the sequence length gives the number of sequences
+        return len(self.indexed_tokens) // self.seq_length
+
+    def __getitem__(self, idx):
+        # Slice the indexed_tokens to get a sequence
+        start_idx = idx * self.seq_length
+        end_idx = start_idx + self.seq_length + 1  # +1 for <eos> token
+        sequence = self.indexed_tokens[start_idx:end_idx]
+        # Convert to torch tensor
+        return torch.tensor(sequence, dtype=torch.long)
+
+# Load the text from a file
+with open('/Users/sean/Downloads/tiny-shakespeare.txt', 'r', encoding='utf-8') as file:
+    text = file.read()
+
+# Create the dataset
+dataset = TextDataset(text)
+#vocab_size = len(list(set(dataset.tokens)))
+vocab_size = len(dataset.vocab)  # Update vocab_size for the transformer model
+train_data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+# Function to generate padding mask for sequences
+def create_padding_mask(seq):
+    return (seq == dataset.vocab['<pad>']).transpose(0, 1)
+
+# Training loop remains the same as before

main.py

@@ -0,0 +1,9 @@
+# List of file paths
+file_paths = ["load_data.py", "transformer_model.py", "train.py"]
+
+# Loop through the file paths
+for file_path in file_paths:
+    # Use exec() to run the code from each file
+    with open(file_path, 'r') as file:
+        code = file.read()
+        exec(code)

predict.py

@@ -0,0 +1,118 @@
+import torch
+import json
+from transformer_model import TransformerModel
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# # Hyperparameters
+d_model = 512  # Dimension of the embeddings and the token representations
+seq_length = 10  # Length of the input and output sequences
+vocab_size = 25672  # Size of the vocabulary
+batch_size = 32  # Batch size for training
+num_heads = 8  # Number of heads in multi-head attention
+dim_feedforward = 2048  # Dimension of feedforward network in encoder and decoder
+
+# Assuming the TransformerModel class is defined in the script
+model = TransformerModel(vocab_size, d_model, num_heads, dim_feedforward, seq_length)
+model.load_state_dict(torch.load('transformer_model.pth'))
+model.eval()  # Set the model to evaluation mode
+
+# Load the vocabulary
+with open('vocabulary.json', 'r') as vocab_file:
+    vocab = json.load(vocab_file)
+
+if '<unk>' not in vocab:
+    # Assign the next integer index to <unk>
+    vocab['<unk>'] = len(vocab)
+
+
+
+def text_to_tensor(text, vocab, seq_length):
+    tokens = text.split()
+    indices = [vocab.get(token, vocab['<unk>']) for token in tokens]  # Replace unknown tokens with <unk>
+    indices = indices[:seq_length]
+    indices += [vocab['<pad>']] * (seq_length - len(indices))
+    return torch.tensor(indices, dtype=torch.long).unsqueeze(0)  # Add batch dimension
+
+
+
+input_text = "please make the"
+input_tensor = text_to_tensor(input_text, vocab, seq_length)
+src = input_tensor
+tgt = input_tensor
+
+
+# def generate_square_subsequent_mask(sz):
+#     mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+#     mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+#     return mask
+
+# def create_padding_mask(seq):
+#     return (seq == vocab['<pad>']).transpose(0, 1)
+
+# Function to generate a square subsequent mask
+def generate_square_subsequent_mask(sz):
+    mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1).bool()
+    return mask
+
+# Function to create padding mask
+def create_padding_mask(seq):
+    mask = (seq == 0).transpose(0, 1) # Assuming 0 is the padding index
+    return mask
+
+src_seq_len = src.size(1)
+tgt_seq_len = tgt.size(1)
+
+src_mask = generate_square_subsequent_mask(src_seq_len)
+# src_mask = torch.zeros((src_seq_len, src_seq_len)).type(torch.bool)
+tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
+src_key_padding_mask = create_padding_mask(src)
+tgt_key_padding_mask = create_padding_mask(tgt)
+
+src.size()
+tgt.size()
+src_mask.size()
+tgt_mask.size()
+src_key_padding_mask.size()
+tgt_key_padding_mask.size()
+
+with torch.no_grad():
+    output = model(src, tgt, src_mask, tgt_mask, 
+                   src_key_padding_mask.transpose(0, 1), tgt_key_padding_mask.transpose(0, 1))
+
+
+predicted_indices = torch.argmax(output, dim=-1).squeeze(0).tolist()
+predicted_indices
+
+inverse_vocab = {value: key for key, value in vocab.items()}
+
+import itertools
+
+flattened_list = list(itertools.chain.from_iterable(predicted_indices))
+
+[inverse_vocab[key] for key in flattened_list]
+
+
+def generate_prediction(text, model, vocab, seq_length):
+    model.eval()  # Make sure the model is in eval mode
+
+    # Convert text to tensor
+    input_tensor = text_to_tensor(text, vocab, seq_length)
+
+    # Generate prediction
+    with torch.no_grad():
+        output = model(input_tensor, input_tensor)  # For simplicity, using the same tensor as src and tgt
+
+    # Convert output tensor to tokens (you may need additional post-processing)
+    predicted_indices = torch.argmax(output, dim=-1).squeeze(0).tolist()
+    predicted_tokens = [vocab[index] for index in predicted_indices]
+
+    return predicted_tokens
+
+
+# Example usage
+text = """Here were the servants of your adversary And 
+yours"""
+
+prediction = generate_prediction(text, model, vocab, seq_length)
+print(prediction)

train.py

@@ -0,0 +1,83 @@
+# # Hyperparameters
+# d_model = 512  # Dimension of the embeddings and the token representations
+# seq_length = 10  # Length of the input and output sequences
+# vocab_size = 1000  # Size of the vocabulary
+# batch_size = 32  # Batch size for training
+# num_heads = 8  # Number of heads in multi-head attention
+# dim_feedforward = 2048  # Dimension of feedforward network in encoder and decoder
+# max_seq_length = seq_length # Maximum sequence length
+
+# # Initialize the model
+# transformer_model = TransformerModel(vocab_size, d_model, num_heads, dim_feedforward, max_seq_length)
+transformer_model = transformer_model.to(device)
+
+# Define loss function and optimizer
+loss_fn = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(transformer_model.parameters(), lr=0.001)
+
+# Function to generate a square subsequent mask
+def generate_square_subsequent_mask(sz):
+    mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1).bool()
+    return mask
+
+# Function to create padding mask
+def create_padding_mask(seq):
+    mask = (seq == 0).transpose(0, 1) # Assuming 0 is the padding index
+    return mask
+
+# Training Loop
+num_epochs = 10
+for epoch in range(num_epochs):
+    transformer_model.train()
+    epoch_loss = 0
+
+    for batch_idx, batch in enumerate(train_data_loader):
+        src_batch = batch[:, :-1].to(device)
+        tgt_input = batch[:, :-1].to(device)  # Excludes the <eos> token
+        targets = batch[:, 1:].contiguous().view(-1).to(device)  # Shifted by one for teacher forcing
+
+        optimizer.zero_grad()
+
+        # Generate square subsequent masks and padding masks
+        src_seq_len = src_batch.size(0)#(1)
+        tgt_seq_len = tgt_input.size(0)#(1)
+        src_mask = generate_square_subsequent_mask(src_seq_len)
+        tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
+        src_padding_mask = create_padding_mask(src_batch)
+        tgt_padding_mask = create_padding_mask(tgt_input)
+
+        # Forward pass
+        output = transformer_model(
+            src_batch.transpose(0,1), tgt_input.transpose(0,1),
+            src_mask, tgt_mask,
+            src_padding_mask, tgt_padding_mask
+        )
+
+# src_batch.size() # torch.Size([9, 10]) (A, B)
+# tgt_input.size() # torch.Size([9, 10]) (A, B)
+# src_mask.size() # torch.Size([9, 9]) (A, A)
+# tgt_mask.size() # torch.Size([9, 9]) (A, A)
+# src_padding_mask.size() # torch.Size([10, 9]) (C, A)
+# tgt_padding_mask.size() # torch.Size([10, 9]) (C, A)
+
+        output = output.view(-1, vocab_size)
+        loss = loss_fn(output, targets)
+
+        loss.backward()
+        optimizer.step()
+
+        epoch_loss += loss.item()
+
+        if batch_idx % 100 == 0:  # Print loss every 100 batches
+            print(f"Epoch {epoch+1}/{num_epochs}, Batch {batch_idx}, Loss: {epoch_loss / (batch_idx+1):.4f}")
+
+    print(f"Epoch {epoch+1}/{num_epochs} finished, Average Loss: {epoch_loss / len(train_data_loader):.4f}")
+
+# Save the model and vocab
+torch.save(transformer_model.state_dict(), 'transformer_model.pth')
+
+import json
+
+with open('vocabulary.json', 'w') as vocab_file:
+    json.dump(dataset.vocab, vocab_file)
+

transformer_model.py

@@ -0,0 +1,188 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Hyperparameters
+d_model = 512  # Dimension of the embeddings and the token representations
+# batch_size = 32  # Batch size for training
+num_heads = 8  # Number of heads in multi-head attention
+dim_feedforward = 2048  # Dimension of feedforward network in encoder and decoder
+
+vocab_size = 25672
+seq_length = 10  # Length of the input and output sequences
+
+# To generate text, we need to define several components:
+# 1. Token Embedding: To convert token indices to vectors
+# 2. Positional Encoding: To give the model information about the order of tokens
+# 3. Transformer Block: Consisting of an encoder and a decoder
+# 4. Output Layer: To convert the transformer output to token probabilities
+
+# Token Embedding
+class TokenEmbedding(nn.Module):
+    def __init__(self, vocab_size, d_model):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+    
+    def forward(self, tokens):
+        return self.embedding(tokens)
+
+# max_len=5000
+# pe = torch.zeros(max_len, d_model)
+# position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+# div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
+# pe[:, 0::2] = torch.sin(position * div_term)
+# pe[:, 1::2] = torch.cos(position * div_term)
+# pe = pe.unsqueeze(0).transpose(0, 1)
+# pe = pe[:max_len]
+
+# Positional Encoding
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        self.max_len = max_len
+        self.d_model = d_model
+        self.register_buffer('pe', self.generate_positional_encoding())
+
+    def generate_positional_encoding(self):
+        # Create a long enough 'pe' for 'max_len'
+        pe = torch.zeros(self.max_len, self.d_model)
+        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / self.d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        return pe[:self.max_len]  # Trim to match the desired max_len
+
+    def forward(self, x):
+        # Add positional encoding to the input tensor
+        # print("x: " + str(x.size()))
+        # print("pe: " + str(pe.size()))
+
+        # x is of shape [batch size, seq length, d_model]
+        # Select positional encoding up to the sequence length of x
+        pe = self.pe[:x.size(0), :]  # pe is now of shape [seq length, d_model]
+        #pez = pe[:src.size(0), :]  # pe is now of shape [seq length, d_model]
+        
+        # Add a batch dimension to pe to make it [1, seq length, d_model]
+        # Then use expand to match the batch size of x
+        # pe = pe.unsqueeze(0).expand(x.size(0), -1, -1)  # pe is now of shape [batch size, seq length, d_model]
+        # pe = pe.unsqueeze(0).expand(src.size(0), -1, -1)
+        
+        # Broadcast the positional encoding to the entire batch
+        # Broadcasting automatically adjusts the dimensions to match x
+        x = x + pe
+        #x = src + pez
+        return x
+    
+
+# Transformer Block
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, num_heads, dim_feedforward, dropout=0.1):
+        super().__init__()
+        # Encoder Layer
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=num_heads,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=10)
+        
+        # Decoder Layer
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=d_model,
+            nhead=num_heads,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout
+        )
+        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=10)
+
+    def forward(self, src, tgt, src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask):
+        # Forward pass through the encoder and decoder
+        memory = self.encoder(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
+        #memory = encoder(src, mask=src_mask, src_key_padding_mask=src_padding_mask)
+# >>> src.size()
+# torch.Size([32, 10, 512])
+# >>> src_mask.size()
+# torch.Size([10, 10])
+# >>> src_padding_mask.size()
+# torch.Size([10, 32])
+        output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=None,
+                              tgt_key_padding_mask=tgt_key_padding_mask,
+                              memory_key_padding_mask=src_key_padding_mask)
+        # output = decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=None,
+                            #   tgt_key_padding_mask=tgt_padding_mask,
+                            #   memory_key_padding_mask=src_padding_mask)
+        return output
+
+# Output Layer
+class OutputLayer(nn.Module):
+    def __init__(self, d_model, vocab_size):
+        super().__init__()
+        self.linear = nn.Linear(d_model, vocab_size)
+    
+    def forward(self, x):
+        return F.log_softmax(self.linear(x), dim=-1)
+
+# src_batch, tgt_input,
+#             src_mask, tgt_mask,
+#             src_padding_mask, tgt_padding_mask
+
+#         src = embedding(src_batch) * torch.sqrt(torch.tensor(d_model))
+#         tgt = embedding(tgt_input) * torch.sqrt(torch.tensor(d_model))
+        # src = pos_encoder(src)
+        # tgt = pos_encoder(tgt)
+#         transformer_output = transformer_block(src, tgt, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask)
+#         out = output_layer(transformer_output)
+
+
+# Putting it all together
+class TransformerModel(nn.Module):
+    def __init__(self, vocab_size, d_model, num_heads, dim_feedforward, max_seq_length):
+        super().__init__()
+        self.embedding = TokenEmbedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(d_model, max_seq_length)
+        #pos_encoder = PositionalEncoding(d_model, seq_length)
+        self.transformer_block = TransformerBlock(d_model, num_heads, dim_feedforward)
+        self.output_layer = OutputLayer(d_model, vocab_size)
+    
+    def forward(self, src, tgt, src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask):
+        src = self.embedding(src) * torch.sqrt(torch.tensor(d_model))
+        tgt = self.embedding(tgt) * torch.sqrt(torch.tensor(d_model))
+        src = self.pos_encoder(src)
+        tgt = self.pos_encoder(tgt)
+        # src = self.pos_encoder(src.transpose(0, 1)).transpose(0, 1)
+        # tgt = self.pos_encoder(tgt.transpose(0, 1)).transpose(0, 1)
+        # src = pos_encoder(src.transpose(0, 1))
+        # tgt = pos_encoder(tgt.transpose(0, 1))
+        transformer_output = self.transformer_block(src.transpose(0, 1), tgt.transpose(0, 1), src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask)
+        #transformer_output = transformer_block(src.transpose(0, 1), tgt.transpose(0, 1), src_mask, tgt_mask, src_padding_mask, tgt_padding_mask)
+        out = self.output_layer(transformer_output)
+        return out
+
+# # Create a sample source and target batch
+# src = torch.randint(low=0, high=vocab_size, size=(seq_length, batch_size))
+# tgt = torch.randint(low=0, high=vocab_size, size=(seq_length, batch_size))
+
+# # Masks and Padding
+# src_mask = torch.zeros((seq_length, seq_length)).type(torch.bool)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# tgt_mask = nn.Transformer.generate_square_subsequent_mask(seq_length, device=device)
+# #tgt_mask = nn.Transformer.generate_square_subsequent_mask(None, seq_length)
+
+# src_key_padding_mask = torch.zeros(batch_size, seq_length).type(torch.bool)  # Assuming no padding
+# tgt_key_padding_mask = torch.zeros(batch_size, seq_length).type(torch.bool)  # Assuming no padding
+
+# Initialize the model
+transformer_model = TransformerModel(vocab_size, d_model, num_heads, dim_feedforward, seq_length)
+
+# Forward pass
+#output = transformer_model(src, tgt, src_mask, tgt_mask, src_key_padding_mask, tgt_key_padding_mask)
+
+# Show the output size
+#print(output.size())  # (seq_length, batch_size, vocab_size)
+
+# In this code, we define a simple Transformer model for educational purposes.
+# It can be expanded by increasing the number of layers, adding dropout, and including more complex masking.
+# The actual training loop, loss calculation, and optimization steps are not shown.
+