# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_data.ipynb.

# %% auto 0
__all__ = ['Data']

# %% ../nbs/01_data.ipynb 3
from transformers import AutoTokenizer
from datasets import Dataset
from datasets import load_dataset, concatenate_datasets

class Data:
    def __init__(self, inputs, outputs, model_id="google/flan-t5-base", verbose=1):
        """
        Preprocess a sequence-to-sequence dataset
        """
        self.ins = inputs
        self.outs = outputs
        self.verbose = 1
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.tokenized_inputs, self.tokenized_outputs = self.truncate()
        
    def truncate(self):
        """
        Truncate input and target texts.
        """
        ins, outs = Dataset.from_dict({'text': self.ins}), Dataset.from_dict({'text':self.outs})
        
        
        tokenized_inputs = ins.map(lambda x: self.tokenizer(x["text"], truncation=True), batched=True, remove_columns=["text"])
        max_source_length = max([len(x) for x in tokenized_inputs["input_ids"]])
        if self.verbose: print(f"Max source length: {max_source_length}")

        # The maximum total sequence length for target text after tokenization. 
        # Sequences longer than this will be truncated, sequences shorter will be padded."
        tokenized_targets = outs.map(lambda x: self.tokenizer(x["text"], truncation=True), batched=True, remove_columns=["text"])
        max_target_length = max([len(x) for x in tokenized_targets["input_ids"]])
        if self.verbose: print(f"Max target length: {max_target_length}")
        return tokenized_inputs, tokenized_targets
    
