code: def train_discriminator(
      dataset,
      dataset_fp=None,
      pretrained_model="gpt2-medium",
      epochs=10,
      batch_size=64,
      log_interval=10,
      save_model=False,
      cached=False,
      no_cuda=False,
  ):
      device = (
          "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"
      )
  
      print("Preprocessing {} dataset...".format(dataset))
      start = time.time()
  
      if dataset == "SST":
          idx2class = [
              "positive",
              "negative",
              "very positive",
              "very negative",
              "neutral",
          ]
          class2idx = {c: i for i, c in enumerate(idx2class)}
  
          discriminator = Discriminator(
              class_size=len(idx2class),
              pretrained_model=pretrained_model,
              cached_mode=cached,
              device=device,
          ).to(device)
  
          text = torchtext_data.Field()
          label = torchtext_data.Field(sequential=False)
          train_data, val_data, test_data = datasets.SST.splits(
              text,
              label,
              fine_grained=True,
              train_subtrees=True,
          )
  
          x = []
          y = []
          for i in trange(len(train_data), ascii=True):
              seq = TreebankWordDetokenizer().detokenize(
                  vars(train_data[i])["text"]
              )
              seq = discriminator.tokenizer.encode(seq)
              seq = torch.tensor(
                  [50256] + seq, device=device, dtype=torch.long
              )
              x.append(seq)
              y.append(class2idx[vars(train_data[i])["label"]])
          train_dataset = Dataset(x, y)
  
          test_x = []
          test_y = []
          for i in trange(len(test_data), ascii=True):
              seq = TreebankWordDetokenizer().detokenize(
                  vars(test_data[i])["text"]
              )
              seq = discriminator.tokenizer.encode(seq)
              seq = torch.tensor(
                  [50256] + seq, device=device, dtype=torch.long
              )
              test_x.append(seq)
              test_y.append(class2idx[vars(test_data[i])["label"]])
          test_dataset = Dataset(test_x, test_y)
  
          discriminator_meta = {
              "class_size": len(idx2class),
              "embed_size": discriminator.embed_size,
              "pretrained_model": pretrained_model,
              "class_vocab": class2idx,
              "default_class": 2,
          }
  
      elif dataset == "clickbait":
          idx2class = ["non_clickbait", "clickbait"]
          class2idx = {c: i for i, c in enumerate(idx2class)}
  
          discriminator = Discriminator(
              class_size=len(idx2class),
              pretrained_model=pretrained_model,
              cached_mode=cached,
              device=device,
          ).to(device)
  
          with open(
              "datasets/clickbait/clickbait_train_prefix.txt"
          ) as f:
              data = []
              for i, line in enumerate(f):
                  try:
                      data.append(eval(line))
                  except Exception:
                      print(
                          "Error evaluating line {}: {}".format(i, line)
                      )
                      continue
          x = []
          y = []
          with open(
              "datasets/clickbait/clickbait_train_prefix.txt"
          ) as f:
              for i, line in enumerate(tqdm(f, ascii=True)):
                  try:
                      d = eval(line)
                      seq = discriminator.tokenizer.encode(d["text"])
  
                      if len(seq) < max_length_seq:
                          seq = torch.tensor(
                              [50256] + seq,
                              device=device,
                              dtype=torch.long,
                          )
                      else:
                          print(
                              "Line {} is longer than maximum length {}".format(
                                  i, max_length_seq
                              )
                          )
                          continue
                      x.append(seq)
                      y.append(d["label"])
                  except Exception:
                      print(
                          "Error evaluating / tokenizing line {}, skipping it".format(
                              i
                          )
                      )
                      pass
  
          full_dataset = Dataset(x, y)
          train_size = int(0.9 * len(full_dataset))
          test_size = len(full_dataset) - train_size
          train_dataset, test_dataset = torch.utils.data.random_split(
              full_dataset, [train_size, test_size]
          )
  
          discriminator_meta = {
              "class_size": len(idx2class),
              "embed_size": discriminator.embed_size,
              "pretrained_model": pretrained_model,
              "class_vocab": class2idx,
              "default_class": 1,
          }
  
      elif dataset == "toxic":
          idx2class = ["non_toxic", "toxic"]
          class2idx = {c: i for i, c in enumerate(idx2class)}
  
          discriminator = Discriminator(
              class_size=len(idx2class),
              pretrained_model=pretrained_model,
              cached_mode=cached,
              device=device,
          ).to(device)
  
          x = []
          y = []
          with open("datasets/toxic/toxic_train.txt") as f:
              for i, line in enumerate(tqdm(f, ascii=True)):
                  try:
                      d = eval(line)
                      seq = discriminator.tokenizer.encode(d["text"])
  
                      if len(seq) < max_length_seq:
                          seq = torch.tensor(
                              [50256] + seq,
                              device=device,
                              dtype=torch.long,
                          )
                      else:
                          print(
                              "Line {} is longer than maximum length {}".format(
                                  i, max_length_seq
                              )
                          )
                          continue
                      x.append(seq)
                      y.append(int(np.sum(d["label"]) > 0))
                  except Exception:
                      print(
                          "Error evaluating / tokenizing line {}, skipping it".format(
                              i
                          )
                      )
                      pass
  
          full_dataset = Dataset(x, y)
          train_size = int(0.9 * len(full_dataset))
          test_size = len(full_dataset) - train_size
          train_dataset, test_dataset = torch.utils.data.random_split(
              full_dataset, [train_size, test_size]
          )
  
          discriminator_meta = {
              "class_size": len(idx2class),
              "embed_size": discriminator.embed_size,
              "pretrained_model": pretrained_model,
              "class_vocab": class2idx,
              "default_class": 0,
          }
  
      else:  # if dataset == "generic":
          # This assumes the input dataset is a TSV with the following structure:
          # class \t text
  
          if dataset_fp is None:
              raise ValueError(
                  "When generic dataset is selected, dataset_fp needs to be specified aswell."
              )
  
          classes = set()
          with open(dataset_fp) as f:
              csv_reader = csv.reader(f, delimiter="\t")
              for row in tqdm(csv_reader, ascii=True):
                  if row:
                      classes.add(row[0])
  
          idx2class = sorted(classes)
          class2idx = {c: i for i, c in enumerate(idx2class)}
  
          discriminator = Discriminator(
              class_size=len(idx2class),
              pretrained_model=pretrained_model,
              cached_mode=cached,
              device=device,
          ).to(device)
  
          x = []
          y = []
          with open(dataset_fp) as f:
              csv_reader = csv.reader(f, delimiter="\t")
              for i, row in enumerate(tqdm(csv_reader, ascii=True)):
                  if row:
                      label = row[0]
                      text = row[1]
  
                      try:
                          seq = discriminator.tokenizer.encode(text)
                          if len(seq) < max_length_seq:
                              seq = torch.tensor(
                                  [50256] + seq,
                                  device=device,
                                  dtype=torch.long,
                              )
  
                          else:
                              print(
                                  "Line {} is longer than maximum length {}".format(
                                      i, max_length_seq
                                  )
                              )
                              continue
  
                          x.append(seq)
                          y.append(class2idx[label])
  
                      except Exception:
                          print(
                              "Error tokenizing line {}, skipping it".format(
                                  i
                              )
                          )
                          pass
  
          full_dataset = Dataset(x, y)
          train_size = int(0.9 * len(full_dataset))
          test_size = len(full_dataset) - train_size
          train_dataset, test_dataset = torch.utils.data.random_split(
              full_dataset, [train_size, test_size]
          )
  
          discriminator_meta = {
              "class_size": len(idx2class),
              "embed_size": discriminator.embed_size,
              "pretrained_model": pretrained_model,
              "class_vocab": class2idx,
              "default_class": 0,
          }
  
      end = time.time()
      print(
          "Preprocessed {} data points".format(
              len(train_dataset) + len(test_dataset)
          )
      )
      print("Data preprocessing took: {:.3f}s".format(end - start))
  
      if cached:
          print("Building representation cache...")
  
          start = time.time()
  
          train_loader = get_cached_data_loader(
              train_dataset,
              batch_size,
              discriminator,
              shuffle=True,
              device=device,
          )
  
          test_loader = get_cached_data_loader(
              test_dataset, batch_size, discriminator, device=device
          )
  
          end = time.time()
          print(
              "Building representation cache took: {:.3f}s".format(
                  end - start
              )
          )
  
      else:
          train_loader = torch.utils.data.DataLoader(
              dataset=train_dataset,
              batch_size=batch_size,
              shuffle=True,
              collate_fn=collate_fn,
          )
          test_loader = torch.utils.data.DataLoader(
              dataset=test_dataset,
              batch_size=batch_size,
              collate_fn=collate_fn,
          )
  
      if save_model:
          with open(
              "{}_classifier_head_meta.json".format(dataset), "w"
          ) as meta_file:
              json.dump(discriminator_meta, meta_file)
  
      optimizer = optim.Adam(discriminator.parameters(), lr=0.0001)
  
      for epoch in range(epochs):
          start = time.time()
          print("\nEpoch", epoch + 1)
  
          train_epoch(
              discriminator=discriminator,
              data_loader=train_loader,
              optimizer=optimizer,
              epoch=epoch,
              log_interval=log_interval,
              device=device,
          )
          evaluate_performance(
              data_loader=test_loader,
              discriminator=discriminator,
              device=device,
          )
  
          end = time.time()
          print("Epoch took: {:.3f}s".format(end - start))
  
          print("\nExample prediction")
          predict(
              example_sentence,
              discriminator,
              idx2class,
              cached=cached,
              device=device,
          )
  
          if save_model:
              # torch.save(discriminator.state_dict(),
              #           "{}_discriminator_{}.pt".format(
              #               args.dataset, epoch + 1
              #               ))
              torch.save(
                  discriminator.get_classifier().state_dict(),
                  "{}_classifier_head_epoch_{}.pt".format(
                      dataset, epoch + 1
                  ),
              )
  
imports: import argparse
  import csv
  import json
  import math
  import time
  import numpy
  import torch
  import torch.optim
  import torch.utils.data
  from nltk.tokenize.treebank import TreebankWordDetokenizer
  from torch import nn
  from torchtext import data
  from torchtext import datasets
  from tqdm import tqdm, trange
  from pplm_classification_head import ClassificationHead
  from transformers import GPT2LMHeadModel, GPT2Tokenizer
label_desc: Use image node_id 7 for calling a global function with example usage: train_discriminator(dataset, dataset_fp, pretrained_model, epochs, batch_size, log_interval, save_model, cached, no_cuda) without return types