LangaugeModelTextBatches.py

class DataLM():
  
  def __init__(self, train_data, index=0,bptt=10, bs=64):
      self.train_data = train_data
      self.index = index
      self.bptt = bptt
      self.bs = bs
      
  def __iter__(self):
    return self
  
  def __next__(self):
    input_seq = []
    output_seq = []
    
    if self.index + (self.bptt * self.bs) + 1 >= len(self.train_data):
      raise StopIteration
    
    index = self.index
    
    #make a batch
    for _ in range(self.bs):
      input_line = []
      output_line = []
      
      #make a line
      for _ in range(self.bptt):
        input_line.append(train_list[index])
        #output line is 1 step ahead of input line
        output_line.append(train_list[index + 1])
        index += 1
      input_seq.append(input_line)
      output_seq.append(output_line)
    self.index = index
    
    for i in range(len(input_seq)):
      #need to map string to int
      input_seq[i] = [string2int[string] for string in input_seq[i]]
      output_seq[i] = [string2int[string] for string in output_seq[i]]
    #convert arrays into tensors and place on gpu
    input_seq = torch.tensor(input_seq, dtype=torch.long).cuda()
    output_seq = torch.tensor(output_seq, dtype=torch.long).cuda()
    return input_seq, output_seq