docstr: Custom Dataset - Converts the input text and label to tensor
  :param tokenizer: bert tokenizer
  :param source: data source - Either a dataframe or DataPipe
  :param max_length: maximum length of the news text
  :param num_samples: number of samples to load
  :param dataset: Dataset type - 20newsgroups or ag_news
docstr_tok: ["Custom","Dataset","-","Converts","the","input","text","and","label","to","tensor",":","param","tokenizer",":","bert","tokenizer",":","param","source",":","data","source","-","Either","a","dataframe","or","DataPipe",":","param","max_length",":","maximum","length","of","the","news","text",":","param","num_samples",":","number","of","samples","to","load",":","param","dataset",":","Dataset","type","-","20newsgroups","or","ag_news"]
code: def __init__(
      self,
      tokenizer,
      source,
      max_length,
      num_samples,
      dataset="20newsgroups",
  ):
      """
      Custom Dataset - Converts the input text and label to tensor
      :param tokenizer: bert tokenizer
      :param source: data source - Either a dataframe or DataPipe
      :param max_length: maximum length of the news text
      :param num_samples: number of samples to load
      :param dataset: Dataset type - 20newsgroups or ag_news
      """
      super().__init__()
      self.source = source
      self.start = 0
      self.tokenizer = tokenizer
      self.max_length = max_length
      self.dataset = dataset
      self.end = num_samples
  
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 1 to create a new NewsDataset object from inherited base classes: IterDataPipe with example: obj = NewsDataset(tokenizer, source, max_length, num_samples, dataset)