<jupyter_start><jupyter_text>text.data.seq2seq.translation> This module contains the bits required to use the fastai DataBlock API and/or mid-level data processing pipelines to organize your data for translation tasks<jupyter_code># export
from typing import Optional

import numpy as np
import pandas as pd

from datasets import Dataset
from fastai.data.block import DataBlock
from transformers import AutoModelForSeq2SeqLM, PreTrainedTokenizerBase, logging

from blurr.text.data.seq2seq.core import Seq2SeqBatchTokenizeTransform, Seq2SeqPreprocessor, Seq2SeqTextBlock
from blurr.text.utils import get_hf_objects

logging.set_verbosity_error()

# hide_input
import os, ast, pdb
from functools import reduce

from datasets import load_dataset
from fastai.data.transforms import *
from fastai.torch_core import *
from fastai.torch_imports import *
from fastcore.all import *
from fastcore.test import *
from nbdev.showdoc import show_doc

from blurr.utils import print_versions
from blurr.text.utils import BlurrText

NLP = BlurrText()

os.environ["TOKENIZERS_PARALLELISM"] = "false"
print("What we're running with at the time this documentation was generated:")
print_versions("torch fastai transformers")

# hide
# cuda
torch.cuda.set_device(1)
print(f"Using GPU #{torch.cuda.current_device()}: {torch.cuda.get_device_name()}")<jupyter_output>Using GPU #1: GeForce GTX 1080 Ti<jupyter_text>SetupWe'll use a subset of `wmt16` to demonstrate how to configure your BLURR for translation tasks<jupyter_code>raw_dataset = load_dataset("wmt16", "de-en", split="train[:1%]")
raw_dataset

print(raw_dataset[0].keys())
print(raw_dataset[0])

wmt_df = pd.DataFrame(raw_dataset["translation"], columns=["de", "en"])

print(len(wmt_df))
wmt_df.head(2)

pretrained_model_name = "facebook/bart-large-cnn"
model_cls = AutoModelForSeq2SeqLM

hf_arch, hf_config, hf_tokenizer, hf_model = get_hf_objects(pretrained_model_name, model_cls=model_cls)
hf_arch, type(hf_tokenizer), type(hf_config), type(hf_model)<jupyter_output><empty_output><jupyter_text>PreprocessingStarting with version 2.0, BLURR provides a preprocessing base class that can be used to build task specific pre-processed datasets from pandas DataFrames or Hugging Face Datasets  `TranslationPreprocessor` -<jupyter_code># export
class TranslationPreprocessor(Seq2SeqPreprocessor):
    def __init__(
        self,
        # A Hugging Face tokenizer
        hf_tokenizer: PreTrainedTokenizerBase,
        # The number of examples to process at a time
        batch_size: int = 1000,
        # The unique identifier in the dataset
        id_attr: Optional[str] = None,
        # The attribute holding the text to translate
        text_attr: str = "original_text",
        # The maximum length (# of tokens) allowed for inputs. Will default to the max length allowed
        # by the model if not provided
        max_input_tok_length: Optional[int] = None,
        # The attribute holding the summary
        target_text_attr: str = "translated_text",
        # The maximum length (# of tokens) allowed for targets
        max_target_tok_length: Optional[int] = None,
        # The attribute that should be created if your are processing individual training and validation
        # datasets into a single dataset, and will indicate to which each example is associated
        is_valid_attr: Optional[str] = "is_valid",
        # Tokenization kwargs that will be applied with calling the tokenizer
        tok_kwargs: dict = {},
    ):
        # we need to use the offset mappings to get back at the raw text from its tokenized representation
        tok_kwargs = {**tok_kwargs, "return_offsets_mapping": True}

        super().__init__(
            hf_tokenizer, batch_size, text_attr, max_input_tok_length, target_text_attr, max_target_tok_length, is_valid_attr, tok_kwargs
        )

        self.id_attr = id_attr

    def process_df(self, training_df: pd.DataFrame, validation_df: Optional[pd.DataFrame] = None):
        df = super().process_df(training_df, validation_df)

        # process df in mini-batches
        final_df = pd.DataFrame()
        for g, batch_df in df.groupby(np.arange(len(df)) // self.batch_size):
            final_df = final_df.append(self._process_df_batch(batch_df))

        final_df.reset_index(drop=True, inplace=True)
        return final_df

    def process_hf_dataset(self, training_ds: Dataset, validation_ds: Optional[Dataset] = None):
        ds = super().process_hf_dataset(training_ds, validation_ds)
        return Dataset.from_pandas(self.process_df(pd.DataFrame(ds)))

    # ----- utility methods -----
    def _process_df_batch(self, batch_df):
        batch_df.reset_index(drop=True, inplace=True)

        # grab our inputs and targets batch encoding objects
        inputs, targets = self._tokenize_function(batch_df.to_dict(orient="list"))

        # add are processed text and target texts to the batched DataFrame
        for txt_seq_idx, (txt_attr, batch_enc) in enumerate(zip([self.text_attr, self.target_text_attr], [inputs, targets])):
            if txt_attr is None:
                break

            char_idxs = []
            for idx, offset_mapping in enumerate(batch_enc["offset_mapping"]):
                text_offsets = [offset_mapping[i] for i, seq_id in enumerate(batch_enc.sequence_ids(idx))]
                char_idxs.append([min(text_offsets)[0], max(text_offsets)[1]])

            batch_df = pd.concat(
                [batch_df, pd.DataFrame(char_idxs, columns=[f"{txt_attr}_start_char_idx", f"{txt_attr}_end_char_idx"])], axis=1
            )
            batch_df.insert(
                0,
                f"proc_{txt_attr}",
                batch_df.apply(lambda r: r[txt_attr][r[f"{txt_attr}_start_char_idx"] : r[f"{txt_attr}_end_char_idx"] + 1], axis=1),
            )

        return batch_df<jupyter_output><empty_output><jupyter_text>This class can be used for preprocessing translation tasks, and includes a `proc_{your_text_attr}` and `proc_{target_text_attr}` attributes containing your modified input and target texts as a result of tokenization (e.g., if you specify a `max_length` the `proc_{your_text_attr}` may contain truncated text).   Using a `DataFrame`<jupyter_code>preprocessor = TranslationPreprocessor(
    hf_tokenizer, text_attr="de", target_text_attr="en", max_input_tok_length=128, max_target_tok_length=128
)
proc_df = preprocessor.process_df(wmt_df)
proc_df.columns, len(proc_df)
proc_df.head(2)<jupyter_output><empty_output><jupyter_text>Examples  Using the mid-level API  Batch-Time Tokenization  Step 1: Get your Hugging Face objects.<jupyter_code>pretrained_model_name = "facebook/bart-large-cnn"
model_cls = AutoModelForSeq2SeqLM

hf_arch, hf_config, hf_tokenizer, hf_model = get_hf_objects(pretrained_model_name, model_cls=model_cls)<jupyter_output><empty_output><jupyter_text>Step 2: Create your `DataBlock`Two lines!  Notice we pass in `noop` for our targets (e.g. our summaries) because the batch transform will take care of both out inputs and targets.<jupyter_code>blocks = (Seq2SeqTextBlock(hf_arch, hf_config, hf_tokenizer, hf_model), noop)
dblock = DataBlock(blocks=blocks, get_x=ColReader("de"), get_y=ColReader("en"), splitter=RandomSplitter())

# dblock.summary(wmt_df)<jupyter_output><empty_output><jupyter_text>Step 3: Build your `DataLoaders`<jupyter_code>dls = dblock.dataloaders(wmt_df, bs=4)

b = dls.one_batch()

len(b), b[0]["input_ids"].shape, b[0]["labels"].shape, b[1].shape

b[0]["labels"][0], b[1][0]

dls.show_batch(dataloaders=dls, max_n=2, input_trunc_at=250, target_trunc_at=250)<jupyter_output><empty_output><jupyter_text>Using a preprocessed dataset  Step 1a: Get your Hugging Face objects.<jupyter_code>pretrained_model_name = "facebook/bart-large-cnn"
model_cls = AutoModelForSeq2SeqLM

hf_arch, hf_config, hf_tokenizer, hf_model = get_hf_objects(pretrained_model_name, model_cls=model_cls)<jupyter_output><empty_output><jupyter_text>Step 1b. Preprocess dataset<jupyter_code>preprocessor = TranslationPreprocessor(
    hf_tokenizer, text_attr="de", target_text_attr="en", max_input_tok_length=128, max_target_tok_length=128
)
proc_df = preprocessor.process_df(wmt_df)<jupyter_output><empty_output><jupyter_text>Step 2: Create your `DataBlock`<jupyter_code>blocks = (Seq2SeqTextBlock(hf_arch, hf_config, hf_tokenizer, hf_model), noop)
dblock = DataBlock(blocks=blocks, get_x=ColReader("proc_de"), get_y=ColReader("proc_en"), splitter=RandomSplitter())<jupyter_output><empty_output><jupyter_text>Step 3: Build your `DataLoaders`<jupyter_code>dls = dblock.dataloaders(proc_df, bs=4)

b = dls.one_batch()

len(b), b[0]["input_ids"].shape, b[0]["labels"].shape, b[1].shape

dls.show_batch(dataloaders=dls, max_n=2, input_trunc_at=250, target_trunc_at=250)<jupyter_output><empty_output><jupyter_text>TestsThe purpose of the following tests is to ensure as much as possible, that the core DataBlock code above works for the pretrained **translation models** below.  These tests are excluded from the CI workflow because of how long they would take to run and the amount of data that would be required to download.**Note**: Feel free to modify the code below to test whatever pretrained translation models you are working with ... and if any of your pretrained summarization models fail, please submit a github issue *(or a PR if you'd like to fix it yourself)*<jupyter_code>[model_type for model_type in NLP.get_models(task="ConditionalGeneration") if (not model_type.startswith("TF"))]

pretrained_model_names = [
    "facebook/bart-base",
    "facebook/wmt19-de-en",  # FSMT
    "Helsinki-NLP/opus-mt-de-en",  # MarianMT
    "sshleifer/tiny-mbart",
    "google/mt5-small",
    "t5-small",
]

path = Path("./")
wmt_df = pd.DataFrame(raw_dataset["translation"], columns=["de", "en"])

# slow
# hide_output
model_cls = AutoModelForSeq2SeqLM
bsz = 2
seq_sz = 128
trg_seq_sz = 128

test_results = []
for model_name in pretrained_model_names:
    error = None

    print(f"=== {model_name} ===\n")

    hf_tok_kwargs = {}
    if model_name == "sshleifer/tiny-mbart":
        hf_tok_kwargs["src_lang"], hf_tok_kwargs["tgt_lang"] = "de_DE", "en_XX"

    hf_arch, hf_config, hf_tokenizer, hf_model = get_hf_objects(model_name, model_cls=model_cls, tokenizer_kwargs=hf_tok_kwargs)

    print(f"architecture:\t{hf_arch}\ntokenizer:\t{type(hf_tokenizer).__name__}\n")

    # not all architectures include a native pad_token (e.g., gpt2, ctrl, etc...), so we add one here
    if hf_tokenizer.pad_token is None:
        hf_tokenizer.add_special_tokens({"pad_token": "<pad>"})
        hf_config.pad_token_id = hf_tokenizer.get_vocab()["<pad>"]
        hf_model.resize_token_embeddings(len(hf_tokenizer))

    batch_tokenize_tfm = Seq2SeqBatchTokenizeTransform(
        hf_arch, hf_config, hf_tokenizer, hf_model, padding="max_length", max_length=seq_sz, max_target_length=trg_seq_sz
    )

    def add_t5_prefix(inp):
        return f"translate German to English: {inp}" if (hf_arch == "t5") else inp

    blocks = (Seq2SeqTextBlock(batch_tokenize_tfm=batch_tokenize_tfm), noop)
    dblock = DataBlock(blocks=blocks, get_x=Pipeline([ColReader("de"), add_t5_prefix]), get_y=ColReader("en"), splitter=RandomSplitter())

    dls = dblock.dataloaders(wmt_df, bs=bsz)
    b = dls.one_batch()

    try:
        print("*** TESTING DataLoaders ***\n")
        test_eq(len(b), 2)
        test_eq(len(b[0]["input_ids"]), bsz)
        test_eq(b[0]["input_ids"].shape, torch.Size([bsz, seq_sz]))
        test_eq(len(b[1]), bsz)
        test_eq(b[1].shape, torch.Size([bsz, trg_seq_sz]))

        if hasattr(hf_tokenizer, "add_prefix_space"):
            test_eq(hf_tokenizer.add_prefix_space, True)

        test_results.append((hf_arch, type(hf_tokenizer).__name__, model_name, "PASSED", ""))
        dls.show_batch(dataloaders=dls, max_n=2, input_trunc_at=1000)

    except Exception as err:
        test_results.append((hf_arch, type(hf_tokenizer).__name__, model_name, "FAILED", err))

# slow
# hide_input
test_results_df = pd.DataFrame(test_results, columns=["arch", "tokenizer", "model_name", "result", "error"])
display_df(test_results_df)<jupyter_output><empty_output><jupyter_text>Export -<jupyter_code># hide
from nbdev.export import notebook2script

notebook2script()<jupyter_output>Converted 00_callbacks.ipynb.
Converted 00_utils.ipynb.
Converted 01_text-callbacks.ipynb.
Converted 01_text-utils.ipynb.
Converted 11_text-data-core.ipynb.
Converted 11_text-modeling-core.ipynb.
Converted 12_text-data-language-modeling.ipynb.
Converted 12_text-modeling-language-modeling.ipynb.
Converted 13_text-data-token-classification.ipynb.
Converted 13_text-modeling-token-classification.ipynb.
Converted 14_text-data-question-answering.ipynb.
Converted 14_text-modeling-question-answering.ipynb.
Converted 20_text-data-seq2seq-core.ipynb.
Converted 20_text-modeling-seq2seq-core.ipynb.
Converted 21_text-data-seq2seq-summarization.ipynb.
Converted 21_text-modeling-seq2seq-summarization.ipynb.
Converted 22_text-data-seq2seq-translation.ipynb.
Converted 22_text-modeling-seq2seq-translation.ipynb.
Converted 99a_text-examples-high-level-api.ipynb.
Converted 99b_text-examples-glue.ipynb.
Converted 99c_text-examples-glue-plain-pytorch.ipynb.
Converted 99d_text-examples-multilabel.ipynb.
Conv[...]
{'fraction_non_alphanumeric': 0.106184012066365, 'fraction_numerical': 0.012368024132730015, 'mean_word_length': 4.885929871282734, 'pattern_counts': {'":': 3, '<': 53, '<?xml version=': 0, '>': 54, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '13', 'dir': 'jupyter-structured-clean-dedup', 'id': '130606', 'license': 'Apache-2.0', 'n_tokens_mistral': 4621, 'n_tokens_neox': 4476, 'n_words': 986, 'path': 'nbs/22_text-data-seq2seq-translation.ipynb', 'repo_name': 'warner-benjamin/blurr', 'size': '13260'}