Answers.Body: <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>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-04-16T22:16:40.877Z
Answers.LastActivityDate: 2023-04-16T22:16:40.877Z
Answers.UserDisplayName: JB1
Body: <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>
  
Tags: <deep-learning><transformer><pytorch><large-language-models>
Title: Playing around with Transformer - accuracy not improving
UserDisplayName: JB1
ThreadText: # 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:
  
  1. 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).
  2. 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]))
  
  ```
  
  
  ---
  
  