class_docstr: Dataset for training
  
  Parameters
  ----------
  dataset : KGDataset
      Original dataset.
  args :
      Global configs.
  ranks:
      Number of partitions.
class_docstr_tok: ["Dataset","for","training","Parameters","--","--","--","--","--","dataset",":","KGDataset","Original","dataset",".","args",":","Global","configs",".","ranks",":","Number","of","partitions","."]
docstr: Create sampler for training
  
  Parameters
  ----------
  batch_size : int
      Batch size of each mini batch.
  neg_sample_size : int
      How many negative edges sampled for each node.
  neg_chunk_size : int
      How many edges in one chunk. We split one batch into chunks.
  mode : str
      Sampling mode.
  number_workers: int
      Number of workers used in parallel for this sampler
  shuffle : bool
      If True, shuffle the seed edges.
      If False, do not shuffle the seed edges.
      Default: False
  exclude_positive : bool
      If True, exlucde true positive edges in sampled negative edges
      If False, return all sampled negative edges even there are positive edges
      Default: False
  rank : int
      Which partition to sample.
  
  Returns
  -------
  dgl.contrib.sampling.EdgeSampler
      Edge sampler
docstr_tok: ["Create","sampler","for","training","Parameters","--","--","--","--","--","batch_size",":","int","Batch","size","of","each","mini","batch",".","neg_sample_size",":","int","How","many","negative","edges","sampled","for","each","node",".","neg_chunk_size",":","int","How","many","edges","in","one","chunk",".","We","split","one","batch","into","chunks",".","mode",":","str","Sampling","mode",".","number_workers",":","int","Number","of","workers","used","in","parallel","for","this","sampler","shuffle",":","bool","If","True",",","shuffle","the","seed","edges",".","If","False",",","do","not","shuffle","the","seed","edges",".","Default",":","False","exclude_positive",":","bool","If","True",",","exlucde","true","positive","edges","in","sampled","negative","edges","If","False",",","return","all","sampled","negative","edges","even","there","are","positive","edges","Default",":","False","rank",":","int","Which","partition","to","sample",".","Returns","--","--","--","-","dgl.contrib.sampling.EdgeSampler","Edge","sampler"]
code: def create_sampler(
      self,
      batch_size,
      neg_sample_size=2,
      neg_chunk_size=None,
      mode="head",
      num_workers=32,
      shuffle=True,
      exclude_positive=False,
      rank=0,
  ):
      """Create sampler for training
  
      Parameters
      ----------
      batch_size : int
          Batch size of each mini batch.
      neg_sample_size : int
          How many negative edges sampled for each node.
      neg_chunk_size : int
          How many edges in one chunk. We split one batch into chunks.
      mode : str
          Sampling mode.
      number_workers: int
          Number of workers used in parallel for this sampler
      shuffle : bool
          If True, shuffle the seed edges.
          If False, do not shuffle the seed edges.
          Default: False
      exclude_positive : bool
          If True, exlucde true positive edges in sampled negative edges
          If False, return all sampled negative edges even there are positive edges
          Default: False
      rank : int
          Which partition to sample.
  
      Returns
      -------
      dgl.contrib.sampling.EdgeSampler
          Edge sampler
      """
      EdgeSampler = getattr(dgl.contrib.sampling, "EdgeSampler")
      assert (
          batch_size % neg_sample_size == 0
      ), "batch_size should be divisible by B"
      return EdgeSampler(
          self.g,
          seed_edges=F.tensor(self.edge_parts[rank]),
          batch_size=batch_size,
          neg_sample_size=int(neg_sample_size / neg_chunk_size),
          chunk_size=neg_chunk_size,
          negative_mode=mode,
          num_workers=num_workers,
          shuffle=shuffle,
          exclude_positive=exclude_positive,
          return_false_neg=False,
      )
  
imports: import math
  import numpy
  import scipy
  import dgl.backend
  import dgl
  import os
  import sys
  import pickle
  import time
  from dgl.base import NID, EID
  import pdb
label_desc: Use image node_id 2 for calling the TrainDataset obj's underlying member method code with example usage: obj.create_sampler(batch_size, neg_sample_size, neg_chunk_size, mode, num_workers, shuffle, exclude_positive, rank) and returns: EdgeSampler