docstr: Split the data into train, test, validation data
  
  :param stage: Stage - training or testing
docstr_tok: ["Split","the","data","into","train",",","test",",","validation","data",":","param","stage",":","Stage","-","training","or","testing"]
code: def setup(self, stage=None):
      """
      Split the data into train, test, validation data
  
      :param stage: Stage - training or testing
      """
      self.tokenizer = BertTokenizer.from_pretrained(
          self.PRE_TRAINED_MODEL_NAME
      )
      if self.dataset == "20newsgroups":
          num_samples = self.num_samples
          self.news_group_df = (
              get_20newsgroups(num_samples)
              if self.dataset == "20newsgroups"
              else get_ag_news(num_samples)
          )
      else:
          train_iter, test_iter = AG_NEWS()
          self.train_dataset = to_map_style_dataset(train_iter)
          self.test_dataset = to_map_style_dataset(test_iter)
  
      if stage == "fit":
          if self.dataset == "20newsgroups":
              self.train_dataset, self.test_dataset = train_test_split(
                  self.news_group_df,
                  test_size=0.3,
                  random_state=self.RANDOM_SEED,
                  stratify=self.news_group_df["label"],
              )
              self.val_dataset, self.test_dataset = train_test_split(
                  self.test_dataset,
                  test_size=0.5,
                  random_state=self.RANDOM_SEED,
                  stratify=self.test_dataset["label"],
              )
  
              self.train_count = len(self.train_dataset)
              self.val_count = len(self.val_dataset)
              self.test_count = len(self.test_dataset)
          else:
              num_train = int(len(self.train_dataset) * 0.95)
              self.train_dataset, self.val_dataset = random_split(
                  self.train_dataset,
                  [num_train, len(self.train_dataset) - num_train],
              )
  
              self.train_count = self.num_samples
              self.val_count = int(self.train_count / 10)
              self.test_count = int(self.train_count / 10)
              self.train_count = self.train_count - (
                  self.val_count + self.test_count
              )
  
          print(
              f"Number of samples used for training: {self.train_count}"
          )
          print(
              f"Number of samples used for validation: {self.val_count}"
          )
          print(f"Number of samples used for test: {self.test_count}")
  
imports: import math
  import os
  import lightning
  import pandas
  import torch
  import torch.nn.functional
  import torchtext.datasets
  from lightning.pytorch.callbacks import EarlyStopping, LearningRateMonitor, ModelCheckpoint
  from lightning.pytorch.cli import LightningCLI
  from sklearn.datasets import fetch_20newsgroups
  from sklearn.metrics import accuracy_score
  from sklearn.model_selection import train_test_split
  from torch import nn
  from torch.utils.data import DataLoader
  from torch.utils.data.dataset import random_split
  from torchdata.datapipes.iter import IterDataPipe
  from torchtext.data.functional import to_map_style_dataset
  from torchtext.datasets import AG_NEWS
  from transformers import BertModel, BertTokenizer
  import mlflow.pytorch
label_desc: Use image node_id 2 for calling the BertDataModule obj's underlying member method code with example usage: obj.setup(stage) without return types