application_vnd.jupyter.stderr_output_1.png
from sklearn.model_selection import train_test_split

from torch.utils.data import Dataset

from transformers import AdamW, T5Tokenizer, T5ForConditionalGeneration

import pandas as pd

import time

import torch

import torch_xla.core.xla_model as xm

special_tokens_dict = {'pad_token': '<pad>', 'bos_token': '<bos>', 'eos_token': '<eos>', 'additional_special_tokens': ['<PAGESTART>', '<PAGEEND>', '<SECTIONSTART>', '<SECTIONEND>', '<TABLESTART>', '<TABLEEND>', '<CELLSTART>', '<CELLEND>', '<COLHEADERSTART>', '<COLHEADEREND>', '<ROWHEADERSTART>', '<ROWHEADEREND>']}
num_added_toks = config.TOKENIZER.add_special_tokens(special_tokens_dict)

df = pd.read_csv(config.TRAINING_FILE)
train_df, val_df = train_test_split(df, test_size=0.1)
train_df = train_df.reset_index(drop=True)
val_df = val_df.reset_index(drop=True)

class tottoDataset(Dataset):
    def __init__(self, df, tokenizer):
        self.sentence = df['sentence']
        self.table = df['table']
        self.tokenizer = tokenizer
    def __len__(self):
        return len(self.sentence)
    def __getitem__(self, idx):
        inp = (self.table[idx] + '</s>').replace('<page_title>', '<PAGESTART>').replace('</page_title>', '<PAGEEND>').replace('<section_title>', '<SECTIONSTART>').replace('</section_title>', '<SECTIONEND>').replace('<table>', '<TABLESTART>').replace('</table>', '<TABLEEND>').replace('<cell>', '<CELLSTART>').replace('</cell>', '<CELLEND>').replace('<col_header>', '<COLHEADERSTART>').replace('</col_header>', '<COLHEADEREND>').replace('<row_header>', '<ROWHEADERSTART>').replace('</row_header>', '<ROWHEADEREND>')
        out = self.sentence[idx] + '</s>'
        inp_tokens = self.tokenizer.encode_plus(inp, padding='max_length', max_length=config.MAX_LEN_I, truncation=True)
        out_tokens = self.tokenizer.encode_plus(out, padding='max_length', max_length=config.MAX_LEN_O, truncation=True)
        inp_id = inp_tokens.input_ids
        out_id = out_tokens.input_ids
        inp_mask = inp_tokens.attention_mask
        out_mask = out_tokens.attention_mask
        labels = out_tokens.input_ids.copy()
        labels = [-100 if x == self.tokenizer.pad_token_id else x for x in labels]
        return {'table_text': inp, 'sentence': out, 'input_ids': torch.tensor(inp_id, dtype=torch.long), 'input_attention_mask': torch.tensor(inp_mask, dtype=torch.long), 'decoder_input_ids': torch.tensor(out_id, dtype=torch.long), 'decoder_attention_mask': torch.tensor(out_mask, dtype=torch.long), 'labels': torch.tensor(labels, dtype=torch.long)}

def train_fn(dataloader, model, optimizer, device, scheduler, epoch, num_epoch, num_steps):
    model.train()
    for i, batch in enumerate(dataloader):
        input_ids = batch['input_ids'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss
        loss.backward()
        xm.optimizer_step(optimizer)
        if scheduler is not None:
            scheduler.step()

def eval_fn(dataloader, model, device, size):
    model.eval()
    loss = 0
    with torch.no_grad():
        for i, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)
            outputs = model(input_ids=input_ids, labels=labels)
            loss += outputs.loss.item()
    return loss

model = T5ForConditionalGeneration.from_pretrained('t5-base', return_dict=True)
model.encoder.resize_token_embeddings(len(config.TOKENIZER))
model.decoder.resize_token_embeddings(len(config.TOKENIZER))