# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['memDataset', 'memData']

# %% ../nbs/00_core.ipynb 3
import torch
from dotmap import DotMap
from sklearn.model_selection import train_test_split
import numpy as np

# %% ../nbs/00_core.ipynb 4
class memDataset(torch.utils.data.Dataset):
    "This dataset holds all your text data in memory and has methods to serve it to a torch data loader."

    def __init__(self, texts: list[str], # List of texts with all your data
                 labels: list[int] = None, # List of labels corresponding to texts
                 name: str = None, # Name of the dataset (e.g. train, test, all, holdout)
                 train=True # Train or Inference type dataset
                 ):
        self.texts = texts
        self.labels = labels
        self.train = train
        self.encodings = None
        self.is_tokenised = False
        self.name = name

    def tokenize(self, 
                 tokenizer, # A tokenizer from Hugging face transformers 
                 **kwargs # Arguments to pass to the tokenizer
                 ):
        self.encodings = tokenizer(self.texts, **kwargs)
        self.is_tokenised = True

    def __getitem__(self, idx):
        if not self.is_tokenised:
            raise TypeError(
                "We need to tokenize the dataset before being able to select elements")
        item = {key: torch.tensor(val[idx])
                for key, val in self.encodings.items()}
        if self.train:
            item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

    def __str__(self) -> str:
        out = [f"name: {self.name}"]
        out.append(f"texts: {self.texts[0]}, ...")
        if self.train:
            out.append(f"labels: {self.labels[0]}, ...")
        out.append(f"tokenised: {self.is_tokenised}")

        return '\n'.join(out)

    def __repr__(self) -> str:
        return self.__str__()


# %% ../nbs/00_core.ipynb 11
class memData():
    "We hold several memDatasets in this class to make it easier to deal with splitting, holdout data or adding augmentations"

    def __init__(self,
                 texts: list[str] = None, # List of all texts available for training
                 labels: list[str] = None, # List of all corresponding labels
                 tokenizer=None, # Tokenizer to be used in the task from `transformers` library
                 test_size=0.2, # Size of the test-train split
                 seed=42) -> None: # Seed to be used in the split
        self.data = DotMap()
        self.test_size = test_size
        self.seed = seed
        self.tokenizer = tokenizer
        if None not in (texts, labels):
            all_data = memDataset(texts, labels, "all")
            self.add_data(all_data)
            self.split(self.test_size)

    def add_data(self, dataset: memDataset):
        self.data[dataset.name] = dataset

    def split(self, test_size):
        indices = np.arange(len(self.data.all))
        train_texts, test_texts, train_labels, test_labels, train_indices, test_indices = train_test_split(
            self.data.all.texts, self.data.all.labels, indices,
            test_size=test_size,
            stratify=self.data.all.labels,
            random_state=self.seed)
        self.add_data(memDataset(train_texts, train_labels, "train"))
        self.add_data(memDataset(test_texts, test_labels, "test"))

    def tokenize_datasets(self, **kwargs):
        if self.tokenizer is None:
            raise ValueError("Tokenizer has not been defined")
            # Any data dict gets tokenized
        for k in self.data.keys():
            if not self.data[k].is_tokenized:
                self.data[k].tokenize(self.tokenizer, **kwargs)

