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:
  
  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]))
  
  
  ---
  
  