import torch
from datasets import load_dataset
from torch.utils.data import Dataset


class NetflixDataset(Dataset):
    def __init__(self, tokenizer):
        super().__init__()

        self.tokenizer = tokenizer
        self.input_ids = []
        self.attn_masks = []
        self.labels = []
        self.txt_list = netflix_descriptions = load_dataset("hugginglearners/netflix-shows", split="train")[
            "description"
        ]
        self.max_length = max([len(self.tokenizer.encode(description)) for description in netflix_descriptions])

        for txt in self.txt_list:
            encodings_dict = self.tokenizer(
                "</s>" + txt + "</s>", truncation=True, max_length=self.max_length, padding="max_length"
            )
            self.input_ids.append(torch.tensor(encodings_dict["input_ids"]))
            self.attn_masks.append(torch.tensor(encodings_dict["attention_mask"]))

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

    def __getitem__(self, idx):
        return self.input_ids[idx], self.attn_masks[idx]


def netflix_collator(data):
    return {
        "input_ids": torch.stack([x[0] for x in data]),
        "attention_mask": torch.stack([x[1] for x in data]),
        "labels": torch.stack([x[0] for x in data]),
    }