<p>Ok, so the solution to this was that I was randomizing both the key (a-z) and the value (0-9). That was too hard for a simple model to learn, so only randomizing the value worked. I also changed the positional embedding to be learned, which improved learning speed.</p>
<pre><code>import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import random
import string
import time
import math

if torch.cuda.is_available(): 
    dev = &quot;cuda:0&quot; 
else: 
    dev = &quot;cpu&quot; 
device = torch.device(dev) 


def flatten(l):
    return [item for sublist in l for item in sublist]

def random_string(length):
    a = string.ascii_lowercase[0:length]
    d = random.sample(string.digits, length)
    r = ' '.join([a[i] + d[i] for i in range(length)])
    n = random.randint(0, length)-1
    return r + &quot; &quot; + a[n] + d[n]

# Synthetic Dataset
class RandomStringDataset(Dataset):
    def __init__(self, num_samples, seq_length):
        self.num_samples = num_samples
        self.seq_length = seq_length
        self.data = [random_string(seq_length) for _ in range(num_samples)]

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        input_seq = self.data[idx][:-1]
        target_seq = self.data[idx][1:]
        return input_seq, target_seq

# Tokenizer and detokenizer functions
def tokenize(text):
    return [char for char in text]

def detokenize(tokens):
    return ''.join(tokens)

# Map characters to indices and vice versa
vocab = string.ascii_lowercase + string.digits + &quot; &quot;
char_to_idx = {char: idx for idx, char in enumerate(vocab)}
idx_to_char = {idx: char for idx, char in enumerate(vocab)}


# Convert tokens to tensor
def tokens_to_tensor(tokens):
    indices = [char_to_idx[token] for token in tokens]
    return torch.tensor(indices, device=device)


# Convert tensor to tokens
def tensor_to_tokens(tensor):
    return [idx_to_char[idx.item()] for idx in tensor]

def collate_fn(batch):
    inputs, targets = zip(*batch)
    input_tensors = [tokens_to_tensor(seq) for seq in inputs]
    target_tensors = [tokens_to_tensor(seq) for seq in targets]
    return torch.stack(input_tensors), torch.stack(target_tensors)

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(DecoderOnlyTransformer, self).__init__()
        self.pe = nn.Embedding(32, d_model)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model, nhead, batch_first=True),
            num_layers
        )
        self.ln = nn.LayerNorm(d_model)
        self.fc = nn.Linear(d_model, vocab_size)
        
    def forward(self, x):
        B, T = x.shape
        x = self.embedding(x) + self.pe(torch.arange(T, device=device))
        output = self.transformer_decoder(x, x)
        output = self.ln(output)
        output = self.fc(output)
        return output
    
# Hyperparameters

seq_length = 10
batch_size = 16
num_samples = batch_size*1000
num_epochs = 10
d_model = 32
nhead = 4
num_layers = 2
vocab_size = len(vocab)

# Create dataset
dataset = RandomStringDataset(num_samples, seq_length)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

# Initialize model, loss, and optimizer
model = DecoderOnlyTransformer(vocab_size, d_model, nhead, num_layers)
model.to(device)

criterion = nn.CrossEntropyLoss()
lr = 5  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

for epoch in range(num_epochs):
    correct = 0
    count = 0
    total_loss = 0
    start_time = 0
    log_interval = 100
    num_batches = int(num_samples/batch_size)
    for inputs, targets in train_loader:
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.reshape(-1, len(vocab)), targets.flatten())

        probabilities = torch.softmax(outputs[0, -1], dim=-1)
        next_token_idx = torch.argmax(probabilities).item()
        token = idx_to_char[next_token_idx]

        if idx_to_char[targets[0, -1].item()] == token:
            correct+=1
        count+=1
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()
        
        total_loss += loss.item()
        if count % log_interval == 0 and count &gt; 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| epoch {epoch:3d} | {count:5d}/{num_batches:5d} batches | '
                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '
                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f} | '
                  f'Acc: {correct/count}'
                 )
            total_loss = 0
            start_time = time.time()    
    print(f'Epoch: {epoch + 1}/{num_epochs}, Loss: {loss.item()}  Count: {count}')
    
    
# Some evaluation
rs = random_string(10)
print(rs)
t = torch.argmax(torch.softmax(model(tokens_to_tensor(tokenize(rs))[:-1].reshape(1, -1)), dim=-1), dim=-1)
&quot;&quot;.join(tensor_to_tokens(t[0]))
</code></pre>

CC BY-SA 4.0
2023-04-16T22:16:40.877Z
2023-04-16T22:16:40.877Z
JB1
<p>I am playing around with a decoder only transformer model,</p>
<p>The Colab is here if you find that easier
<a href="https://colab.research.google.com/drive/1SHyJ9Oa3E4j1x8YFlXQbd1mjUjWhHGOV#scrollTo=60e13119" rel="nofollow noreferrer">https://colab.research.google.com/drive/1SHyJ9Oa3E4j1x8YFlXQbd1mjUjWhHGOV#scrollTo=60e13119</a>
or see the code below (should run with minimal deps in a notebook)</p>
<h3>Goal:</h3>
<ol>
<li>I am playing around to understand the transformer architecture. I wanted to build a minimal version of a decoder only model and see if i can train the model to predict the last digit in this sequence (think of it as a recall function from a key value store).</li>
<li>Want to get an intuitive understanding on the linear transformations of the weight matrices to the sequence input + positional encodings.</li>
</ol>
<pre><code>n9 v8 a5 p7 k0 j1 i3 e2 g6 c4 c - (should be 4)
d5 t3 q8 r7 y1 i0 c2 n9 s4 u6 i - (should be 0) ...
l8 u9 p5 y3 f7 k0 g6 v4 r1 x2 l
a7 x5 b6 v0 i1 f3 z9 d4 y2 k8 x
m9 h4 g5 t2 l3 f1 w7 b6 a8 j0 g
x0 g7 q9 u2 j8 v4 h3 o1 f5 r6 r
r6 c4 d0 p3 j2 g9 a7 n1 e8 l5 d
r2 z7 y6 x5 v4 u1 s3 a8 l9 p0 z
k0 u3 t1 r4 g8 p2 j5 x9 s7 v6 t
o7 a1 u3 r2 k6 j0 m8 y9 e4 c5 j
</code></pre>
<h3>Questions:</h3>
<ul>
<li>The accuracy is not improving, so i guess there is some fundamental issue with the model. It looks like it is learning that it is a digit, but not what digit.</li>
</ul>
<h3>Code:</h3>
<pre><code>import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import random
import string
import time
import math

if torch.cuda.is_available(): 
    dev = &quot;cuda:0&quot; 
else: 
    dev = &quot;cpu&quot; 
device = torch.device(dev) 

class PositionalEncoding(nn.Module):
    &quot;&quot;&quot;
    compute sinusoid encoding.
    &quot;&quot;&quot;
    def __init__(self, d_model, max_len, device):
        &quot;&quot;&quot;
        constructor of sinusoid encoding class

        :param d_model: dimension of model
        :param max_len: max sequence length
        :param device: hardware device setting
        &quot;&quot;&quot;
        super(PositionalEncoding, self).__init__()

        # same size with input matrix (for adding with input matrix)
        self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False  # we don't need to compute gradient

        pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        # 1D =&gt; 2D unsqueeze to represent word's position

        _2i = torch.arange(0, d_model, step=2, device=device).float()
        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])
        # &quot;step=2&quot; means 'i' multiplied with two (same with 2 * i)

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        # compute positional encoding to consider positional information of words

    def forward(self, x):
        # self.encoding
        # [max_len = 512, d_model = 512]

        batch_size, seq_len, d_model = x.size()
        #print(&quot;seq_len: &quot;, seq_len, &quot; batch_size: &quot;, batch_size, &quot;d_model: &quot;, d_model)
        # [batch_size = 128, seq_len = 31]
        #print(&quot;self.encoding: &quot;, self.encoding[:seq_len, :].shape)
        return self.encoding[:seq_len, :d_model]
        # [seq_len = 30, d_model = 512]
        # it will add with tok_emb : [128, 30, 512]   

# Generate random strings
def random_string(length):
    #return &quot;1 2 3 4 5 6 7 8 9&quot;
    a = random.sample(string.ascii_lowercase, length)
    d = random.sample(string.digits, length)
    r = ' '.join([a[i] + d[i] for i in range(length)])
    n = random.randint(0, length)-1
    return r + &quot; &quot; + a[n] + d[n]

# Synthetic Dataset
class RandomStringDataset(Dataset):
    def __init__(self, num_samples, seq_length):
        self.num_samples = num_samples
        self.seq_length = seq_length
        self.data = [random_string(seq_length) for _ in range(num_samples)]

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        input_seq = self.data[idx][:-1]
        target_seq = self.data[idx][1:]
        return input_seq, target_seq

# Tokenizer and detokenizer functions
def tokenize(text):
    return [char for char in text]

def detokenize(tokens):
    return ''.join(tokens)

# Map characters to indices and vice versa
vocab = string.ascii_lowercase + string.digits + &quot; &quot;
char_to_idx = {char: idx for idx, char in enumerate(vocab)}
idx_to_char = {idx: char for idx, char in enumerate(vocab)}


# Convert tokens to tensor
def tokens_to_tensor(tokens):
    indices = [char_to_idx[token] for token in tokens]
    return torch.tensor(indices, device=device)

# Convert tensor to tokens
def tensor_to_tokens(tensor):
    return [idx_to_char[idx.item()] for idx in tensor]

def collate_fn(batch):
    inputs, targets = zip(*batch)
    input_tensors = [tokens_to_tensor(seq) for seq in inputs]
    target_tensors = [tokens_to_tensor(seq) for seq in targets]
    return torch.stack(input_tensors), torch.stack(target_tensors)

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(DecoderOnlyTransformer, self).__init__()
        self.pe = PositionalEncoding(d_model, 128, device)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model, nhead),
            num_layers
        )
        self.fc = nn.Linear(d_model, vocab_size)
        
    def forward(self, x):
        x = self.embedding(x)
        pe = self.pe(x)
        #print(&quot;x: &quot;, x.shape)
        #print(&quot;pe: &quot;, pe.shape)
        x = x + pe
        tgt = torch.zeros_like(x)
        output = self.transformer_decoder(x, x)
        output = self.fc(output)
        return output
    
# Hyperparameters

seq_length = 10
batch_size = 16
num_samples = batch_size*1000
learning_rate = 0.001
num_epochs = 100
d_model = 4
nhead = 4
num_layers = 2
vocab_size = len(vocab)

# Create dataset
dataset = RandomStringDataset(num_samples, seq_length)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

# Initialize model, loss, and optimizer
model = DecoderOnlyTransformer(vocab_size, d_model, nhead, num_layers)
model.to(device)


criterion = nn.CrossEntropyLoss()
lr = 5  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

for epoch in range(num_epochs):
    correct = 0
    count = 0
    total_loss = 0
    start_time = 0
    log_interval = 100
    num_batches = int(num_samples/batch_size)
    for inputs, targets in train_loader:
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.reshape(-1, len(vocab)), targets.flatten())

        probabilities = torch.softmax(outputs[0, -1], dim=-1)
        next_token_idx = torch.argmax(probabilities).item()
        token = idx_to_char[next_token_idx]

        if idx_to_char[targets[0, -1].item()] == token:
            correct+=1
        count+=1
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()
        
        total_loss += loss.item()
        if count % log_interval == 0 and count &gt; 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| epoch {epoch:3d} | {count:5d}/{num_batches:5d} batches | '
                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '
                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f} | '
                  f'Acc: {correct/count}'
                 )
            total_loss = 0
            start_time = time.time()    
    print(f'Epoch: {epoch + 1}/{num_epochs}, Loss: {loss.item()}  Count: {count}')
<span class="math-container">```</span>
</code></pre>

<deep-learning><transformer><pytorch><large-language-models>
Playing around with Transformer - accuracy not improving
JB1
# Playing around with Transformer - accuracy not improving

(Asked by: username_0 on 2023-04-16)

I am playing around with a decoder only transformer model,
The Colab is here if you find that easier
https://colab.research.google.com/drive/1SHyJ9Oa3E4j1x8YFlXQbd1mjUjWhHGOV#scrollTo=60e13119
or see the code below (should run with minimal deps in a notebook)
Goal:

I am playing around to understand the transformer architecture. I wanted to build a minimal version of a decoder only model and see if i can train the model to predict the last digit in this sequence (think of it as a recall function from a key value store).
Want to get an intuitive understanding on the linear transformations of the weight matrices to the sequence input + positional encodings.

n9 v8 a5 p7 k0 j1 i3 e2 g6 c4 c - (should be 4)
d5 t3 q8 r7 y1 i0 c2 n9 s4 u6 i - (should be 0) ...
l8 u9 p5 y3 f7 k0 g6 v4 r1 x2 l
a7 x5 b6 v0 i1 f3 z9 d4 y2 k8 x
m9 h4 g5 t2 l3 f1 w7 b6 a8 j0 g
x0 g7 q9 u2 j8 v4 h3 o1 f5 r6 r
r6 c4 d0 p3 j2 g9 a7 n1 e8 l5 d
r2 z7 y6 x5 v4 u1 s3 a8 l9 p0 z
k0 u3 t1 r4 g8 p2 j5 x9 s7 v6 t
o7 a1 u3 r2 k6 j0 m8 y9 e4 c5 j

Questions:

The accuracy is not improving, so i guess there is some fundamental issue with the model. It looks like it is learning that it is a digit, but not what digit.

Code:
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import random
import string
import time
import math

if torch.cuda.is_available(): 
    dev = "cuda:0" 
else: 
    dev = "cpu" 
device = torch.device(dev) 

class PositionalEncoding(nn.Module):
    """
    compute sinusoid encoding.
    """
    def __init__(self, d_model, max_len, device):
        """
        constructor of sinusoid encoding class

        :param d_model: dimension of model
        :param max_len: max sequence length
        :param device: hardware device setting
        """
        super(PositionalEncoding, self).__init__()

        # same size with input matrix (for adding with input matrix)
        self.encoding = torch.zeros(max_len, d_model, device=device)
        self.encoding.requires_grad = False  # we don't need to compute gradient

        pos = torch.arange(0, max_len, device=device)
        pos = pos.float().unsqueeze(dim=1)
        # 1D => 2D unsqueeze to represent word's position

        _2i = torch.arange(0, d_model, step=2, device=device).float()
        # 'i' means index of d_model (e.g. embedding size = 50, 'i' = [0,50])
        # "step=2" means 'i' multiplied with two (same with 2 * i)

        self.encoding[:, 0::2] = torch.sin(pos / (10000 ** (_2i / d_model)))
        self.encoding[:, 1::2] = torch.cos(pos / (10000 ** (_2i / d_model)))
        # compute positional encoding to consider positional information of words

    def forward(self, x):
        # self.encoding
        # [max_len = 512, d_model = 512]

        batch_size, seq_len, d_model = x.size()
        #print("seq_len: ", seq_len, " batch_size: ", batch_size, "d_model: ", d_model)
        # [batch_size = 128, seq_len = 31]
        #print("self.encoding: ", self.encoding[:seq_len, :].shape)
        return self.encoding[:seq_len, :d_model]
        # [seq_len = 30, d_model = 512]
        # it will add with tok_emb : [128, 30, 512]   

# Generate random strings
def random_string(length):
    #return "1 2 3 4 5 6 7 8 9"
    a = random.sample(string.ascii_lowercase, length)
    d = random.sample(string.digits, length)
    r = ' '.join([a[i] + d[i] for i in range(length)])
    n = random.randint(0, length)-1
    return r + " " + a[n] + d[n]

# Synthetic Dataset
class RandomStringDataset(Dataset):
    def __init__(self, num_samples, seq_length):
        self.num_samples = num_samples
        self.seq_length = seq_length
        self.data = [random_string(seq_length) for _ in range(num_samples)]

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        input_seq = self.data[idx][:-1]
        target_seq = self.data[idx][1:]
        return input_seq, target_seq

# Tokenizer and detokenizer functions
def tokenize(text):
    return [char for char in text]

def detokenize(tokens):
    return ''.join(tokens)

# Map characters to indices and vice versa
vocab = string.ascii_lowercase + string.digits + " "
char_to_idx = {char: idx for idx, char in enumerate(vocab)}
idx_to_char = {idx: char for idx, char in enumerate(vocab)}


# Convert tokens to tensor
def tokens_to_tensor(tokens):
    indices = [char_to_idx[token] for token in tokens]
    return torch.tensor(indices, device=device)

# Convert tensor to tokens
def tensor_to_tokens(tensor):
    return [idx_to_char[idx.item()] for idx in tensor]

def collate_fn(batch):
    inputs, targets = zip(*batch)
    input_tensors = [tokens_to_tensor(seq) for seq in inputs]
    target_tensors = [tokens_to_tensor(seq) for seq in targets]
    return torch.stack(input_tensors), torch.stack(target_tensors)

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(DecoderOnlyTransformer, self).__init__()
        self.pe = PositionalEncoding(d_model, 128, device)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model, nhead),
            num_layers
        )
        self.fc = nn.Linear(d_model, vocab_size)
        
    def forward(self, x):
        x = self.embedding(x)
        pe = self.pe(x)
        #print("x: ", x.shape)
        #print("pe: ", pe.shape)
        x = x + pe
        tgt = torch.zeros_like(x)
        output = self.transformer_decoder(x, x)
        output = self.fc(output)
        return output
    
# Hyperparameters

seq_length = 10
batch_size = 16
num_samples = batch_size*1000
learning_rate = 0.001
num_epochs = 100
d_model = 4
nhead = 4
num_layers = 2
vocab_size = len(vocab)

# Create dataset
dataset = RandomStringDataset(num_samples, seq_length)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

# Initialize model, loss, and optimizer
model = DecoderOnlyTransformer(vocab_size, d_model, nhead, num_layers)
model.to(device)


criterion = nn.CrossEntropyLoss()
lr = 5  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

for epoch in range(num_epochs):
    correct = 0
    count = 0
    total_loss = 0
    start_time = 0
    log_interval = 100
    num_batches = int(num_samples/batch_size)
    for inputs, targets in train_loader:
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.reshape(-1, len(vocab)), targets.flatten())

        probabilities = torch.softmax(outputs[0, -1], dim=-1)
        next_token_idx = torch.argmax(probabilities).item()
        token = idx_to_char[next_token_idx]

        if idx_to_char[targets[0, -1].item()] == token:
            correct+=1
        count+=1
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()
        
        total_loss += loss.item()
        if count % log_interval == 0 and count > 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| epoch {epoch:3d} | {count:5d}/{num_batches:5d} batches | '
                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '
                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f} | '
                  f'Acc: {correct/count}'
                 )
            total_loss = 0
            start_time = time.time()    
    print(f'Epoch: {epoch + 1}/{num_epochs}, Loss: {loss.item()}  Count: {count}')
```

**Tags:** deep-learning, transformer, pytorch, large-language-models


---

## 1 answers

### Answer by username_0 on 2023-04-16. Score: 0

Ok, so the solution to this was that I was randomizing both the key (a-z) and the value (0-9). That was too hard for a simple model to learn, so only randomizing the value worked. I also changed the positional embedding to be learned, which improved learning speed.
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import random
import string
import time
import math

if torch.cuda.is_available(): 
    dev = "cuda:0" 
else: 
    dev = "cpu" 
device = torch.device(dev) 


def flatten(l):
    return [item for sublist in l for item in sublist]

def random_string(length):
    a = string.ascii_lowercase[0:length]
    d = random.sample(string.digits, length)
    r = ' '.join([a[i] + d[i] for i in range(length)])
    n = random.randint(0, length)-1
    return r + " " + a[n] + d[n]

# Synthetic Dataset
class RandomStringDataset(Dataset):
    def __init__(self, num_samples, seq_length):
        self.num_samples = num_samples
        self.seq_length = seq_length
        self.data = [random_string(seq_length) for _ in range(num_samples)]

    def __len__(self):
        return self.num_samples

    def __getitem__(self, idx):
        input_seq = self.data[idx][:-1]
        target_seq = self.data[idx][1:]
        return input_seq, target_seq

# Tokenizer and detokenizer functions
def tokenize(text):
    return [char for char in text]

def detokenize(tokens):
    return ''.join(tokens)

# Map characters to indices and vice versa
vocab = string.ascii_lowercase + string.digits + " "
char_to_idx = {char: idx for idx, char in enumerate(vocab)}
idx_to_char = {idx: char for idx, char in enumerate(vocab)}


# Convert tokens to tensor
def tokens_to_tensor(tokens):
    indices = [char_to_idx[token] for token in tokens]
    return torch.tensor(indices, device=device)


# Convert tensor to tokens
def tensor_to_tokens(tensor):
    return [idx_to_char[idx.item()] for idx in tensor]

def collate_fn(batch):
    inputs, targets = zip(*batch)
    input_tensors = [tokens_to_tensor(seq) for seq in inputs]
    target_tensors = [tokens_to_tensor(seq) for seq in targets]
    return torch.stack(input_tensors), torch.stack(target_tensors)

class DecoderOnlyTransformer(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers):
        super(DecoderOnlyTransformer, self).__init__()
        self.pe = nn.Embedding(32, d_model)
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer_decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model, nhead, batch_first=True),
            num_layers
        )
        self.ln = nn.LayerNorm(d_model)
        self.fc = nn.Linear(d_model, vocab_size)
        
    def forward(self, x):
        B, T = x.shape
        x = self.embedding(x) + self.pe(torch.arange(T, device=device))
        output = self.transformer_decoder(x, x)
        output = self.ln(output)
        output = self.fc(output)
        return output
    
# Hyperparameters

seq_length = 10
batch_size = 16
num_samples = batch_size*1000
num_epochs = 10
d_model = 32
nhead = 4
num_layers = 2
vocab_size = len(vocab)

# Create dataset
dataset = RandomStringDataset(num_samples, seq_length)
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

# Initialize model, loss, and optimizer
model = DecoderOnlyTransformer(vocab_size, d_model, nhead, num_layers)
model.to(device)

criterion = nn.CrossEntropyLoss()
lr = 5  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

for epoch in range(num_epochs):
    correct = 0
    count = 0
    total_loss = 0
    start_time = 0
    log_interval = 100
    num_batches = int(num_samples/batch_size)
    for inputs, targets in train_loader:
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.reshape(-1, len(vocab)), targets.flatten())

        probabilities = torch.softmax(outputs[0, -1], dim=-1)
        next_token_idx = torch.argmax(probabilities).item()
        token = idx_to_char[next_token_idx]

        if idx_to_char[targets[0, -1].item()] == token:
            correct+=1
        count+=1
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()
        
        total_loss += loss.item()
        if count % log_interval == 0 and count > 0:
            lr = scheduler.get_last_lr()[0]
            ms_per_batch = (time.time() - start_time) * 1000 / log_interval
            cur_loss = total_loss / log_interval
            ppl = math.exp(cur_loss)
            print(f'| epoch {epoch:3d} | {count:5d}/{num_batches:5d} batches | '
                  f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | '
                  f'loss {cur_loss:5.2f} | ppl {ppl:8.2f} | '
                  f'Acc: {correct/count}'
                 )
            total_loss = 0
            start_time = time.time()    
    print(f'Epoch: {epoch + 1}/{num_epochs}, Loss: {loss.item()}  Count: {count}')
    
    
# Some evaluation
rs = random_string(10)
print(rs)
t = torch.argmax(torch.softmax(model(tokens_to_tensor(tokenize(rs))[:-1].reshape(1, -1)), dim=-1), dim=-1)
"".join(tensor_to_tokens(t[0]))


---

