class_docstr: Simple neural tagging model
  Supports pytorch embedder models, multi-gpu training, KD from teacher models
  
  Args:
      embedder_model: pytorch embedder model (valid nn.Module model)
      word_vocab (Vocabulary): word vocabulary
      labels (List, optional): list of labels. Defaults to None
      use_crf (bool, optional): use CRF a the classifier (instead of Softmax). Defaults to False.
      device (str, optional): device backend. Defatuls to 'cpu'.
      n_gpus (int, optional): number of gpus. Default to 0.
docstr: Train a tagging model
  
  Args:
      train_data_set (DataLoader): train examples dataloader.
          - If distiller object is provided train examples should contain a tuple of
            student/teacher data examples.
      dev_data_set (DataLoader, optional): dev examples dataloader. Defaults to None.
      test_data_set (DataLoader, optional): test examples dataloader. Defaults to None.
      epochs (int, optional): num of epochs to train. Defaults to 3.
      batch_size (int, optional): batch size. Defaults to 8.
      optimizer (fn, optional): optimizer function. Defaults to default model optimizer.
      max_grad_norm (float, optional): max gradient norm. Defaults to 5.0.
      logging_steps (int, optional): number of steps between logging. Defaults to 50.
      save_steps (int, optional): number of steps between model saves. Defaults to 100.
      save_path (str, optional): model output path. Defaults to None.
      distiller (TeacherStudentDistill, optional): KD model for training the model using
      a teacher model. Defaults to None.
      best_result_file (str, optional): path to save best dev results when it's updated.
      word_dropout (float, optional): whole-word (-> oov) dropout rate. Defaults to 0.
code: def train(
      self,
      train_data_set: DataLoader,
      dev_data_set: DataLoader = None,
      test_data_set: DataLoader = None,
      epochs: int = 3,
      batch_size: int = 8,
      optimizer=None,
      max_grad_norm: float = 5.0,
      logging_steps: int = 50,
      save_steps: int = 100,
      save_path: str = None,
      distiller: TeacherStudentDistill = None,
      best_result_file: str = None,
      word_dropout: float = 0,
  ):
      """
      Train a tagging model
  
      Args:
          train_data_set (DataLoader): train examples dataloader.
              - If distiller object is provided train examples should contain a tuple of
                student/teacher data examples.
          dev_data_set (DataLoader, optional): dev examples dataloader. Defaults to None.
          test_data_set (DataLoader, optional): test examples dataloader. Defaults to None.
          epochs (int, optional): num of epochs to train. Defaults to 3.
          batch_size (int, optional): batch size. Defaults to 8.
          optimizer (fn, optional): optimizer function. Defaults to default model optimizer.
          max_grad_norm (float, optional): max gradient norm. Defaults to 5.0.
          logging_steps (int, optional): number of steps between logging. Defaults to 50.
          save_steps (int, optional): number of steps between model saves. Defaults to 100.
          save_path (str, optional): model output path. Defaults to None.
          distiller (TeacherStudentDistill, optional): KD model for training the model using
          a teacher model. Defaults to None.
          best_result_file (str, optional): path to save best dev results when it's updated.
          word_dropout (float, optional): whole-word (-> oov) dropout rate. Defaults to 0.
      """
      if optimizer is None:
          optimizer = self.get_optimizer()
      train_batch_size = batch_size * max(1, self.n_gpus)
      logger.info("***** Running training *****")
      logger.info("  Num examples = %d", len(train_data_set.dataset))
      logger.info("  Num Epochs = %d", epochs)
      logger.info(
          "  Instantaneous batch size per GPU/CPU = %d", batch_size
      )
      logger.info("  Total batch size = %d", train_batch_size)
      global_step = 0
      best_dev = 0
      dev_test = 0
      self.model.zero_grad()
      epoch_it = trange(epochs, desc="Epoch")
      for epoch in epoch_it:
          step_it = tqdm(train_data_set, desc="Train iteration")
          avg_loss = 0
  
          for step, batches in enumerate(step_it):
              self.model.train()
  
              batch, t_batch = (
                  (batches, []) if not distiller else (batches[:2])
              )
              batch = tuple(t.to(self.device) for t in batch)
              inputs = self.batch_mapper(batch)
              logits = self.model(**inputs)
  
              if distiller:
                  t_batch = tuple(t.to(self.device) for t in t_batch)
                  t_logits = distiller.get_teacher_logits(t_batch)
                  valid_positions = (
                      t_batch[3] != 0.0
                  )  # TODO: implement method to get only valid logits from the model itself
                  valid_t_logits = {}
                  max_seq_len = logits.shape[1]
                  for i in range(len(logits)):  # each example in batch
                      valid_logit_i = t_logits[i][valid_positions[i]]
                      valid_t_logits[i] = (
                          valid_logit_i
                          if valid_logit_i.shape[0] <= max_seq_len
                          else valid_logit_i[:][:max_seq_len]
                      )  # cut to max len
  
                  # prepare teacher labels for non-labeled examples
                  t_labels_dict = {}
                  for i in range(len(valid_t_logits.keys())):
                      t_labels_dict[i] = torch.argmax(
                          F.log_softmax(valid_t_logits[i], dim=-1),
                          dim=-1,
                      )
  
              # pseudo labeling
              for i, is_labeled in enumerate(inputs["is_labeled"]):
                  if not is_labeled:
                      t_labels_i = t_labels_dict[i]
                      # add the padded teacher label:
                      inputs["labels"][i] = torch.cat(
                          (
                              t_labels_i,
                              torch.zeros(
                                  [max_seq_len - len(t_labels_i)],
                                  dtype=torch.long,
                              ).to(self.device),
                          ),
                          0,
                      )
  
              # apply word dropout to the input
              if word_dropout != 0:
                  tokens = inputs["words"]
                  tokens = np.array(tokens.detach().cpu())
                  word_probs = np.random.random(tokens.shape)
                  drop_indices = np.where(
                      (word_probs > word_dropout) & (tokens != 0)
                  )  # ignore padding indices
                  inputs["words"][
                      drop_indices[0], drop_indices[1]
                  ] = self.word_vocab.oov_id
  
              # loss
              if self.use_crf:
                  loss = -1.0 * self.crf(
                      logits,
                      inputs["labels"],
                      mask=inputs["mask"] != 0.0,
                  )
              else:
                  loss_fn = CrossEntropyLoss(ignore_index=0)
                  loss = loss_fn(
                      logits.view(-1, self.num_labels),
                      inputs["labels"].view(-1),
                  )
  
              # for idcnn training - add dropout penalty loss
              module = (
                  self.model.module if self.n_gpus > 1 else self.model
              )
              if isinstance(module, IDCNN) and module.drop_penalty != 0:
                  logits_no_drop = self.model(**inputs, no_dropout=True)
                  sub = logits.sub(logits_no_drop)
                  drop_loss = torch.div(torch.sum(torch.pow(sub, 2)), 2)
                  loss += module.drop_penalty * drop_loss
  
              if self.n_gpus > 1:
                  loss = loss.mean()
  
              # add distillation loss if activated
              if distiller:
                  # filter masked student logits (no padding)
                  valid_s_logits = {}
                  valid_s_positions = inputs["mask"] != 0.0
                  for i in range(len(logits)):
                      valid_s_logit_i = logits[i][valid_s_positions[i]]
                      valid_s_logits[i] = valid_s_logit_i
                  loss = distiller.distill_loss_dict(
                      loss, valid_s_logits, valid_t_logits
                  )
  
              loss.backward()
              torch.nn.utils.clip_grad_norm_(
                  self.model.parameters(), max_grad_norm
              )
  
              optimizer.step()
              optimizer.zero_grad()
              global_step += 1
              avg_loss += loss.item()
              if global_step % logging_steps == 0:
                  if step != 0:
                      logger.info(
                          " global_step = %s, average loss = %s",
                          global_step,
                          avg_loss / step,
                      )
                      best_dev, dev_test = self.update_best_model(
                          dev_data_set,
                          test_data_set,
                          best_dev,
                          dev_test,
                          best_result_file,
                          avg_loss / step,
                          epoch,
                          save_path=None,
                      )
              if (
                  save_steps != 0
                  and save_path is not None
                  and global_step % save_steps == 0
              ):
                  self.save_model(save_path)
      self.update_best_model(
          dev_data_set,
          test_data_set,
          best_dev,
          dev_test,
          best_result_file,
          "end_training",
          "end_training",
          save_path=save_path + "/best_dev",
      )
  
imports: import io
  import logging
  import os
  import pickle
  from typing import List
  import numpy
  import torch
  import torch.optim
  from torch.nn import CrossEntropyLoss
  from torch.nn import functional
  from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
  from tqdm import tqdm, trange
  from nlp_architect.data.sequential_tagging import TokenClsInputExample
  from nlp_architect.models import TrainableModel
  from nlp_architect.nn.torch.layers import CRF
  from nlp_architect.nn.torch.distillation import TeacherStudentDistill
  from nlp_architect.nn.torch.modules.embedders import IDCNN
  from nlp_architect.utils.metrics import tagging
  from nlp_architect.utils.text import Vocabulary, char_to_id
label_desc: Use image node_id 5 for calling the NeuralTagger obj's underlying member method code with example usage: obj.train(train_data_set, dev_data_set, test_data_set, epochs, batch_size, optimizer, max_grad_norm, logging_steps, save_steps, save_path, distiller, best_result_file, word_dropout) without return types