NER: Data

# from python
import random

# from pypi
import numpy

# this project
from neurotic.nlp.named_entity_recognition import NERData, TOKEN

ner = NERData()

# to make the functions pass we need to use their names (initially)
vocab = vocabulary = ner.data.vocabulary
tag_map = tags = ner.data.tags

print(vocabulary["the"])
print(vocabulary[TOKEN.pad])
print(vocabulary["The"])

for tag, index in tags.items():
    print(f" - {tag}: {index}")

print(f"The number of tags is {len(tag_map)}")
print(f"The vocabulary size is {len(vocab):,}")
print(f"The training size is {len(ner.data.data_sets.x_train):,}")
print(f"The validation size is {len(ner.data.data_sets.x_validate):,}")
print("The first training sentence is ")
print(f"'{' '.join(ner.data.raw_data_sets.x_train[0])}'")
print("Its corresponding label is")
print(f" '{' '.join(ner.data.raw_data_sets.y_train[0])}'")

print("The first training encoded sentence is ")
print(f"{ner.data.data_sets.x_train[0]}")
print("Its corresponding encoded label is")
print(f"{ner.data.data_sets.y_train[0]}")

def data_generator(batch_size: int, x: list, y: list, pad: int,
                   shuffle: bool=False, verbose: bool=False):
    """Generate batches of data for training

    Args: 
      batch_size - size of each batch generated
      x - sentences where words are represented as integers
      y - tags associated with the sentences
      pad - number to use as the padding character
      shuffle - Whether to shuffle the data
      verbose - Whether to print information to stdout

    Yields:
     a tuple containing 2 elements:
       X - np.ndarray of dim (batch_size, max_len) of padded sentences
       Y - np.ndarray of dim (batch_size, max_len) of tags associated with the sentences in X
    """    
    # count the number of lines in data_lines
    num_lines = len(x)

    # create an array with the indexes of data_lines that can be shuffled
    lines_index = list(range(num_lines))

    # shuffle the indexes if shuffle is set to True
    if shuffle:
        random.shuffle(lines_index)

    index = 0 # tracks current location in x, y
    while True:
        buffer_x = [0] * batch_size
        buffer_y = [0] * batch_size
        max_len = 0
        for i in range(batch_size):
             # if the index is greater than or equal to the number of lines in x
            if index >= num_lines:
                # then reset the index to 0
                index = 0
                # re-shuffle the indexes if shuffle is set to True
                if shuffle:
                    random.shuffle(lines_index)

            # The current position is obtained using `lines_index[index]`
            # Store the x value at the current position into the buffer_x
            buffer_x[i] = x[lines_index[index]]

            # Store the y value at the current position into the buffer_y
            buffer_y[i] = y[lines_index[index]]

            lenx = len(buffer_x[i])    #length of current x[]
            if lenx > max_len:
                max_len = lenx                   #max_len tracks longest x[]

            # increment index by one
            index += 1


        # create X,Y, NumPy arrays of size (batch_size, max_len) 'full' of pad value
        X = numpy.full((batch_size, max_len), pad)
        Y = numpy.full((batch_size, max_len), pad)

        # copy values from lists to NumPy arrays. Use the buffered values
        for i in range(batch_size):
            # get the example (sentence as a tensor)
            # in `buffer_x` at the `i` index
            x_i = buffer_x[i]

            # similarly, get the example's labels
            # in `buffer_y` at the `i` index
            y_i = buffer_y[i]

            # Walk through each word in x_i
            for j in range(len(x_i)):
                # store the word in x_i at position j into X
                X[i, j] = x_i[j]

                # store the label in y_i at position j into Y
                Y[i, j] = y_i[j]

        if verbose: print("index=", index)
        yield((X,Y))

batch_size = 5
mini_sentences = ner.data.data_sets.x_train[0: 8]
mini_labels = ner.data.data_sets.y_train[0: 8]
dg = data_generator(batch_size, mini_sentences, mini_labels, vocab["<PAD>"], shuffle=False, verbose=True)
X1, Y1 = next(dg)
X2, Y2 = next(dg)
print(Y1.shape, X1.shape, Y2.shape, X2.shape)
print(X1[0][:], "\n", Y1[0][:])

# from python
from typing import List, Tuple
import random

# from pypi
import attr
import numpy

Vectors = List[List[int]]
Batch = Tuple[numpy.ndarray]

@attr.s(auto_attribs=True)
class DataGenerator:
    """A generator of data to train the NER Model

    Args:
     batch_size: how many lines to generate at once
     x: the encoded sentences
     y: the encoded labels 
     padding: encoding to use for padding lines
     shuffle: whether to shuffle the data
     verbose: whether to print messages to stdout
    """
    batch_size: int
    x: Vectors
    y: Vectors
    padding: int
    shuffle: bool=False
    verbose: bool=False
    _batch: iter=None

from neurotic.nlp.named_entity_recognition import DataGenerator

generator = DataGenerator(x=ner.data.data_sets.x_train[0:8],
                          y=ner.data.data_sets.y_train[0: 8],
                          batch_size=5,
                          padding=vocabulary[TOKEN.pad])

X1, Y1 = next(generator)
X2, Y2 = next(generator)
print(Y1.shape, X1.shape, Y2.shape, X2.shape)
print(X1[0][:], "\n", Y1[0][:])