# AUTOGENERATED! DO NOT EDIT! File to edit: 00_core.ipynb (unless otherwise specified).

__all__ = ['NNSplit']

# Cell
import numpy as np
import torch

# Cell
from .utils import text_to_id, DATA_DIRECTORY, CUT_LENGTH, STRIDE, DEVICE
from .tokenizer import Tokenizer

# Cell
class NNSplit(Tokenizer):
    def __init__(self, model_name, stride=STRIDE,
                 cut_length=CUT_LENGTH,
                 data_directory=DATA_DIRECTORY,
                 device=DEVICE):
        self.cut_length = cut_length
        self.stride = stride
        self.device = device

        self.model = torch.jit.load(str(data_directory / f"{model_name}.pt")).to(self.device)

    def split(self, text):
        all_inputs = []
        all_idx = []

        inputs = [text_to_id(x) for x in text]

        while len(inputs) < self.cut_length:
            inputs.append(0)

        n_preds = np.zeros(len(inputs))

        start = 0
        end = -1
        while end != len(inputs):
            end = min(start + self.cut_length, len(inputs))
            start = end - self.cut_length

            idx = slice(start, end)
            all_inputs.append(inputs[idx])
            all_idx.append(idx)

            n_preds[idx] += 1

            start += self.stride

        batched_inputs = torch.tensor(all_inputs, dtype=torch.int64, device=self.device)
        preds = torch.sigmoid(self.model(batched_inputs).detach().cpu()).numpy()

        avg_preds = np.zeros((len(inputs), 2))

        for pred, idx in zip(preds, all_idx):
            avg_preds[idx] += pred

        avg_preds /= n_preds[:, None]

        sentences = []
        tokens = []
        token = ""

        for char, pred in zip(text, avg_preds):
            token += char

            if pred[0] > threshold:
                tokens.append(token)
                token = ""

            if pred[1] > threshold:
                sentences.append(tokens)
                tokens = []

        if len(token) > 0:
            tokens.append(token)

        if len(tokens) > 0:
            sentences.append(tokens)

        return sentences