# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.8.5 64-bit (''fastai'': conda)'
#     language: python
#     name: python38564bitfastaicondad52d12c5a30a4725bf9d3e235cf1271c
# ---

# %load_ext autoreload
# %autoreload 2

# + _cell_guid="b1076dfc-b9ad-4769-8c92-a6c4dae69d19" _uuid="8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
from fastai.vision.all import *
import sklearn.metrics as skm
from tqdm.notebook import tqdm
import sklearn.feature_extraction.text
from transformers import AutoModelForMaskedLM, AutoConfig, BertTokenizer, LineByLineTextDataset

from shopee_utils import *
from train_utils import *
import codecs
from torch.utils.data.dataset import Dataset
# -

BERT_CONFIG_PATH = './indobert-large-p2'
PATH = Path('/home/slex/data/shopee')

train_df = pd.read_csv(PATH/'train.csv')

model = AutoModelForMaskedLM.from_pretrained(BERT_CONFIG_PATH)

tokenizer = BertTokenizer.from_pretrained(BERT_CONFIG_PATH)


class TextDataset(Dataset):
    def __init__(self, tokenizer, lines, block_size):
        batch_encoding = tokenizer(lines, add_special_tokens=True, truncation=True, max_length=block_size)
        self.examples = batch_encoding["input_ids"]
        self.examples = [{"input_ids": torch.tensor(e, dtype=torch.long)} for e in self.examples]
    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]


dataset=TextDataset(tokenizer, train_df.title.to_list(), 128)

tokenizer.decode(dataset[0]['input_ids'])

# +
from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

# +
from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="./LanguageModel",
    overwrite_output_dir=True,
    num_train_epochs=20,
    per_gpu_train_batch_size=32,
    save_steps=10_000,
    save_total_limit=2,
    prediction_loss_only=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset,
)
# -

trainer.train()

trainer.save_model('./finetuned_indobert-large-p2')



StarcoderdataJupyterScriptsDedupFiltered