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

__all__ = ['Textual', 'HFTextBlock']

# Cell
from bs4 import BeautifulSoup
import requests
import logging
from ipywidgets import interact, interact_manual, FileUpload
from pathlib import Path
from forgebox.html import DOM

import numpy as np
import pandas as pd
from typing import Dict, List, Any

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

    def __init__(self, text: str):
        self.text = text.replace("\n", " ").replace("\r", "")

    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):
        res = requests.get(url)
        if res.status_code == 200:
            res.encoding = 'utf-8'
            headers = res.headers
            if "html" in str(headers).lower():
                text = BeautifulSoup(res.text).text
            else:
                text = res.text
            obj = cls(text)
            obj.path = "./text_data.txt"
            with open(obj.path, "w") as f:
                f.write(obj.text)
            return obj
        else:
            raise ConnectionError(f"Error downloading: {url}")

    @classmethod
    def from_path(cls, path: Path):
        """
        Load a textual object from system path
        """
        if Path(path).exists() == False:
            raise FileExistsError(f"Can not find {path}")
        with open(path, ) as f:
            obj = cls(f.read())
            obj.path = path
            return obj

    @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 fucntion
        - 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 arguements 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