# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/41_nlp_data.ipynb (unless otherwise specified).

__all__ = ['PathStr', 'Textual', 'HFTextBlock']

# Cell
import logging
import re
import requests
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import numpy as np
import pandas as pd
from bs4 import BeautifulSoup
from forgebox.html import DOM
from ipywidgets import interact, interact_manual, FileUpload
from unpackai import utils

PathStr = Union[Path, str]


# Cell
class Textual:
    """
    Obtain and manage textual data
    """

    def __init__(self, text: str, path: Optional[Path] = None):
        self.text = re.sub(r"[\r]?\n", " ", text)
        self.path = path

    def __repr__(self) -> str:
        return f"""Text ({len(self.text)} chars), textual(),
    train_path, val_path = textual.create_train_val()"""

    def __call__(self, page_size: int = 1000) -> None:
        """
        Previewing the first 200(or less) pages
        - page_size: character number each page
        """
        logging.info(f"Previewing the first 200 pages")

        @interact
        def show_text(page=(0, min(len(self.text)//page_size-1, 200), 1)):
            display(self.text[page*page_size: (page+1)*page_size])

    @classmethod
    def from_url(cls, url: str) -> 'Textual':
        """Create a Textual object from a URL"""
        return cls.from_path(utils.download(url))

    @classmethod
    def from_path(cls, path: PathStr) -> 'Textual':
        """
        Load a textual object from system path
        """
        path = Path(path)
        if not path.exists():
            raise FileExistsError(f"Cannot find {path}")
        return cls(path.read_text(encoding="utf-8"), path=path)

    @classmethod
    def from_upload(
        cls,
        path: Path = Path("./uploaded_file.txt")
    ):
        """
        Load textual with interactive upload button
        """
        DOM("🗃 Please upload a text file ended in .txt", "h4")()
        my_manual = interact_manual.options(manual_name="Upload")

        @my_manual
        def create_upload(btn_upload=FileUpload(description="Choose File")):
            text = list(btn_upload.values())[-1]['content'].decode()
            with open(path, "w") as f:
                f.write(text)
            return path

        def uploaded():
            result = create_upload.widget.result
            if result is None:
                raise FileExistsError(
                    "You have to upload the txt file first")
            return cls.from_path(result)
        return uploaded

    def create_train_val(
            self,
            valid_ratio=.2,
            train_path="./train_text.txt",
            val_path="./val_text.txt"):
        """
        create 2 files:
        - ./train_text.txt
        - ./val_text.txt
        """
        split = int(len(self.text)*(valid_ratio))
        with open(train_path, "w") as f:
            f.write(self.text[split:])
        with open(val_path, "w") as f:
            f.write(self.text[:split])
        return train_path, val_path

    def show_batch(self, tokenizer, bs:int = 4):
        from torch.utils.data.dataloader import DataLoader
        bunch = self.create_datasets(tokenizer)
        return next(iter(DataLoader(
            bunch['train_dataset'],
            batch_size=bs,
            collate_fn = bunch["data_collator"],
            )))

    def create_datasets(
        self,
        tokenizer,
        valid_ratio: float = .2,
        train_path: str = "./train_text.txt",
        val_path: str = "./val_text.txt",
        block_size: int = 128,
        mlm: bool = False,
    ) -> Dict[str, Any]:
        """
        Create pytorch datasets and collating functions
        - tokenizer: a huggingface tokenizer
        - valid ratio: portion of the valid data,
            compare to the entire dataset
        - train_path: a path saving train text file
        - val_path: a path saving valid text file
        - block_size: max possible length of the sequence
        - mlm, return a masked language modeling collating
            default False
        """
        # split dataset
        train_path, val_path = self.create_train_val(
            valid_ratio=valid_ratio,
            train_path=train_path,
            val_path=val_path,
        )

        from transformers import TextDataset, DataCollatorForLanguageModeling
        train_dataset = TextDataset(
            tokenizer=tokenizer,
            file_path=train_path,
            block_size=block_size)

        test_dataset = TextDataset(
            tokenizer=tokenizer,
            file_path=val_path,
            block_size=block_size)

        data_collator = DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=mlm,
        )
        return dict(
            train_dataset=train_dataset,
            eval_dataset=test_dataset,
            data_collator=data_collator,
        )

    def get_hf_trainer(
        self,
        model,
        tokenizer,
        arguments = None,
        valid_ratio: float = .2,
        train_path: str = "./train_text.txt",
        val_path: str = "./val_text.txt",
        block_size: int = 128,
        mlm: bool = False,
    ):
        """
        Create a hugging face trainer
        """
        from transformers import TrainingArguments, Trainer
        if arguments is None:
            arguments = TrainingArguments(
                output_dir="./write_style",
                overwrite_output_dir=True,  num_train_epochs=3,
                eval_steps = 400, save_steps=800, warmup_steps=600,
                per_device_train_batch_size=24,
                per_device_eval_batch_size=64,
            )
        trainer = Trainer(
            model=model, args=arguments,
            **self.create_datasets(tokenizer,
                valid_ratio=valid_ratio,
                train_path=train_path,
                val_path=val_path,
                block_size=block_size,
                mlm=mlm)
        )
        return trainer

# Cell


def HFTextBlock(tokenizer, **tk_kwargs):
    """
    create Huggingface specialized fastai Block
    tokenizer: PreTrainedTokenizer from huggingface
    tk_kwargs: keyward arguments for tokenizer's ```__call__```
    """
    from fastai.data.block import TransformBlock
    import torch

    def text_2_tensor_collate(data):
        """
        During the usual collation
        Use tokenizer to encode the text by batch
        """
        cols = list(zip(*data))
        result = []
        for col in cols:
            if type(col[0]) == str:
                result.append(
                    tokenizer(list(col), return_tensors="pt", **tk_kwargs)["input_ids"]
                )
            else:
                result.append(torch.stack(list(col), 0))
        return tuple(result)

    def get_hf_text_block():
        return TransformBlock(
            type_tfms=str, dls_kwargs={"create_batch": text_2_tensor_collate}
        )

    return get_hf_text_block
