class_docstr: ## Dataset
  
  This class loads and pre-processes the data.
class_docstr_tok: ["#","#","Dataset","This","class","loads","and","pre-processes","the","data","."]
docstr: `dataset` is a list of numpy arrays of shape [seq_len, 3].
  It is a sequence of strokes, and each stroke is represented by
  3 integers.
  First two are the displacements along x and y ($\Delta x$, $\Delta y$)
  and the last integer represents the state of the pen, $1$ if it's touching
  the paper and $0$ otherwise.
docstr_tok: ["`","dataset","`","is","a","list","of","numpy","arrays","of","shape","[","seq_len",",","3","]",".","It","is","a","sequence","of","strokes",",","and","each","stroke","is","represented","by","3","integers",".","First","two","are","the","displacements","along","x","and","y","(","$","\\Delta","x","$",",","$","\\Delta","y","$",")","and","the","last","integer","represents","the","state","of","the","pen",",","$","1","$","if","it","'s","touching","the","paper","and","$","0","$","otherwise","."]
code: def __init__(
      self,
      dataset: np.array,
      max_seq_length: int,
      scale: Optional[float] = None,
  ):
      """
      `dataset` is a list of numpy arrays of shape [seq_len, 3].
      It is a sequence of strokes, and each stroke is represented by
      3 integers.
      First two are the displacements along x and y ($\Delta x$, $\Delta y$)
      and the last integer represents the state of the pen, $1$ if it's touching
      the paper and $0$ otherwise.
      """
  
      data = []
      # We iterate through each of the sequences and filter
      for seq in dataset:
          # Filter if the length of the sequence of strokes is within our range
          if 10 < len(seq) <= max_seq_length:
              # Clamp $\Delta x$, $\Delta y$ to $[-1000, 1000]$
              seq = np.minimum(seq, 1000)
              seq = np.maximum(seq, -1000)
              # Convert to a floating point array and add to `data`
              seq = np.array(seq, dtype=np.float32)
              data.append(seq)
  
      # We then calculate the scaling factor which is the
      # standard deviation of ($\Delta x$, $\Delta y$) combined.
      # Paper notes that the mean is not adjusted for simplicity,
      # since the mean is anyway close to $0$.
      if scale is None:
          scale = np.std(
              np.concatenate([np.ravel(s[:, 0:2]) for s in data])
          )
      self.scale = scale
  
      # Get the longest sequence length among all sequences
      longest_seq_len = max([len(seq) for seq in data])
  
      # We initialize PyTorch data array with two extra steps for start-of-sequence (sos)
      # and end-of-sequence (eos).
      # Each step is a vector $(\Delta x, \Delta y, p_1, p_2, p_3)$.
      # Only one of $p_1, p_2, p_3$ is $1$ and the others are $0$.
      # They represent *pen down*, *pen up* and *end-of-sequence* in that order.
      # $p_1$ is $1$ if the pen touches the paper in the next step.
      # $p_2$ is $1$ if the pen doesn't touch the paper in the next step.
      # $p_3$ is $1$ if it is the end of the drawing.
      self.data = torch.zeros(
          len(data), longest_seq_len + 2, 5, dtype=torch.float
      )
      # The mask array needs only one extra-step since it is for the outputs of the
      # decoder, which takes in `data[:-1]` and predicts next step.
      self.mask = torch.zeros(len(data), longest_seq_len + 1)
  
      for i, seq in enumerate(data):
          seq = torch.from_numpy(seq)
          len_seq = len(seq)
          # Scale and set $\Delta x, \Delta y$
          self.data[i, 1 : len_seq + 1, :2] = seq[:, :2] / scale
          # $p_1$
          self.data[i, 1 : len_seq + 1, 2] = 1 - seq[:, 2]
          # $p_2$
          self.data[i, 1 : len_seq + 1, 3] = seq[:, 2]
          # $p_3$
          self.data[i, len_seq + 1 :, 4] = 1
          # Mask is on until end of sequence
          self.mask[i, : len_seq + 1] = 1
  
      # Start-of-sequence is $(0, 0, 1, 0, 0)$
      self.data[:, 0, 2] = 1
  
imports: import math
  from typing import Optional, Tuple, Any
  import numpy
  import torch
  import torch.nn
  from matplotlib import pyplot
  from torch import optim
  from torch.utils.data import Dataset, DataLoader
  import einops
  from labml import lab, experiment, tracker, monit
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.module import Module
  from labml_helpers.optimizer import OptimizerConfigs
  from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
label_desc: Use image node_id 1 to create a new StrokesDataset object from inherited base classes: Dataset with example: obj = StrokesDataset(dataset, max_seq_length, scale)