text: import numpy as np
  import torch
  import math
  from tqdm import tqdm
  from torch.cuda import FloatTensor, LongTensor
  
  def iterate_batches(data, batch_size):
      X, y = data
      n_samples = len(X)
  
      indices = np.arange(n_samples)
      np.random.shuffle(indices)
      
      for start in range(0, n_samples, batch_size):
          end = min(start + batch_size, n_samples)
          
          batch_indices = indices[start:end]
          
          max_sent_len = max(len(X[ind]) for ind in batch_indices)
          X_batch = np.zeros((max_sent_len, len(batch_indices)))
          y_batch = np.zeros((max_sent_len, len(batch_indices)))
          
          for batch_ind, sample_ind in enumerate(batch_indices):
              X_batch[:len(X[sample_ind]), batch_ind] = X[sample_ind]
              y_batch[:len(y[sample_ind]), batch_ind] = y[sample_ind]
              
          yield X_batch, y_batch
  
  
  def do_epoch(model, criterion, data, batch_size, optimizer=None, name=None):
      epoch_loss = 0
      correct_count = 0
      sum_count = 0
  
      is_train = not optimizer is None
      name = name or ''
      model.train(is_train)
      
      batches_count = math.ceil(len(data[0]) / batch_size)
      
      with torch.autograd.set_grad_enabled(is_train):
          with tqdm(total=batches_count) as progress_bar:
              for i, (X_batch, y_batch) in enumerate(iterate_batches(data, batch_size)):
                  X_batch, y_batch = LongTensor(X_batch), LongTensor(y_batch)
                  logits = model(X_batch)
  
                  loss = criterion(logits.reshape(logits.shape[0]*logits.shape[1], logits.shape[2]),
                                   y_batch.reshape(y_batch.shape[0]*y_batch.shape[1]))
  
                  epoch_loss += loss.item()
  
                  if optimizer:
                      optimizer.zero_grad()
                      loss.backward()
                      optimizer.step()
                  
                  preds = torch.argmax(logits, dim=-1)
                  mask = (y_batch != 0).float()
                  
                  cur_correct_count, cur_sum_count = ((preds == y_batch).float() * mask).sum().item(), mask.sum().item()
  
                  correct_count += cur_correct_count
                  sum_count += cur_sum_count
  
                  progress_bar.update()
                  progress_bar.set_description('{:>5s} Loss = {:.5f}, Accuracy = {:.2%}'.format(
                      name, loss.item(), cur_correct_count / cur_sum_count)
                  )
                  
              progress_bar.set_description('{:>5s} Loss = {:.5f}, Accuracy = {:.2%}'.format(
                  name, epoch_loss / batches_count, correct_count / sum_count)
              )
  
      return epoch_loss / batches_count, correct_count / sum_count
  
  
  def fit(model, criterion, optimizer, train_data, epochs_count=1, batch_size=32,
          val_data=None, val_batch_size=None):
          
      if not val_data is None and val_batch_size is None:
          val_batch_size = batch_size
          
      for epoch in range(epochs_count):
          name_prefix = '[{} / {}] '.format(epoch + 1, epochs_count)
          train_loss, train_acc = do_epoch(model, criterion, train_data, batch_size, optimizer, name_prefix + 'Train:')
          
          if not val_data is None:
              val_loss, val_acc = do_epoch(model, criterion, val_data, val_batch_size, None, name_prefix + '  Val:')
  