<jupyter_start><jupyter_text>Translate model We are using [this nice dataset](https://github.com/BangBOOM/Classical-Chinese)  Imports<jupyter_code>from forgebox.imports import *
from forgebox.thunder.callbacks import DataFrameMetricsCallback
from gc_utils.env import *
from datasets import load_dataset
# from fastai.text.all import *
from unpackai.nlp import *
from tqdm.notebook import tqdm
import random
import pytorch_lightning as pl
import re

def remove_all_punkt(text):
    """
    Removes all punctuation from Chinese text.

    :param text: text to remove punctuation from
    :return: text with no punctuation
    """
    return re.sub(r'[^\w\s]', '', text)
remove_all_punkt("亳州水军千户胡进等领骑兵渡淝水，逾荆山，与宋兵战，杀获甚众，赏钞币有差。")<jupyter_output><empty_output><jupyter_text>Config<jupyter_code>DATA = Path(sys_loc('DATA')/"nlp"/"zh"/"cc_vs_zh")
TO_CLASSICAL = False<jupyter_output><empty_output><jupyter_text>Download data  Data  Combine data<jupyter_code>all_file = list(DATA.rglob("data/*"))
def open_file_to_lines(file):
    with open(file) as f:
        lines = f.read().splitlines()
    return lines

def pairing_the_file(files,kw):
    pairs = []
    for file in files:
        if kw not in file.name:
            file1 = file
            file2 = f"{file}{kw}"
            pairs.append((file1,file2))
    return pairs
pairs = pairing_the_file(all_file,"翻译")
def open_pairs(pairs):
    chunks = []
    for pair in tqdm(pairs, leave=False):
        file1,file2 = pair
        lines1 = open_file_to_lines(file1)
        lines2 = open_file_to_lines(file2)
        chunks.append(pd.DataFrame({"classical":lines1,"modern":lines2}))
    return pd.concat(chunks).sample(frac=1.).reset_index(drop=True)
data_df = open_pairs(pairs)
df = data_df.rename(
    columns = dict(
        zip(["modern","classical"],
             ["source","target"] if TO_CLASSICAL else ["target","source",]))
)
df.head()<jupyter_output><empty_output><jupyter_text>Loading tokenizer<jupyter_code>from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    AutoModel,
    EncoderDecoderModel
    )

# we find a English parsing encoder, as a pretrained bert is good at understanding english
# BERT is short for Bidirectional **Encoder** Representations from Transformers, which consists fully of encoder blocks
ENCODER_PRETRAINED = "bert-base-chinese"
# we find a Chinese writing model for decoder, as decoder is the part of the model that can write stuff
DECODER_PRETRAINED = "uer/gpt2-chinese-poem"

encoder_tokenizer = AutoTokenizer.from_pretrained(ENCODER_PRETRAINED)

decoder_tokenizer = AutoTokenizer.from_pretrained(
    ENCODER_PRETRAINED # notice we use the BERT's tokenizer here
)<jupyter_output><empty_output><jupyter_text>Pytoch Dataset<jupyter_code>class Seq2Seq(Dataset):
    def __init__(
        self, df, tokenizer, target_tokenizer,
        max_len=128,
        no_punkt:bool = False,
    ):
        """
        no_punkt, do we ramdomly remove punctuation
        from source sentence
        """
        super().__init__()
        self.df = df
        self.tokenizer = tokenizer
        self.target_tokenizer = target_tokenizer
        self.max_len = max_len
        self.no_punkt = no_punkt
        
    def __len__(self, ):
        return len(self.df)

    def __getitem__(self, idx):
        return dict(self.df.iloc[idx])

    def collate(self, batch):
        batch_df = pd.DataFrame(list(batch))
        x, y = batch_df.source, batch_df.target
        # there is a random no punctuation mode
        # for source text
        # as some of the classical text we get
        # might be whole chunk of paragraph without
        # any punctuation
        if self.no_punkt:
            x = list(i if random.random()>.5
                     else remove_all_punkt(i)
                     for i in x)
        else:
            x = list(x)
        x_batch = self.tokenizer(
            x,
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt',
        )
        y_batch = self.target_tokenizer(
            list(y),
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt',
        )
        x_batch['decoder_input_ids'] = y_batch['input_ids']
        x_batch['labels'] = y_batch['input_ids'].clone()
        x_batch['labels'][x_batch['labels'] == self.tokenizer.pad_token_id] = -100
        return x_batch

    def dataloader(self, batch_size, shuffle=True):
        return DataLoader(
            self,
            batch_size=batch_size,
            shuffle=shuffle,
            collate_fn=self.collate,
        )

    def split_train_valid(self, valid_size=0.1):
        split_index = int(len(self) * (1 - valid_size))
        cls = type(self)
        shuffled = self.df.sample(frac=1).reset_index(drop=True)
        train_set = cls(
            shuffled.iloc[:split_index],
            tokenizer=self.tokenizer,
            target_tokenizer=self.target_tokenizer,
            max_len=self.max_len,
            no_punkt=self.no_punkt,
        )
        valid_set = cls(
            shuffled.iloc[split_index:],
            tokenizer=self.tokenizer,
            target_tokenizer=self.target_tokenizer,
            max_len=self.max_len,
            no_punkt=self.no_punkt,
        )
        return train_set, valid_set<jupyter_output><empty_output><jupyter_text>PL datamodule<jupyter_code>class Seq2SeqData(pl.LightningDataModule):
    def __init__(
        self, df,
        tokenizer,
        target_tokenizer,
        batch_size=12,
        max_len=128,
        no_punkt:bool=False):
        super().__init__()
        self.df = df
        self.ds = Seq2Seq(df,
                          tokenizer,
                          target_tokenizer,
                          max_len=max_len,
                          no_punkt=no_punkt)
        self.tokenizer = tokenizer
        self.target_tokenizer = target_tokenizer
        self.max_len = max_len
        self.batch_size = batch_size

    def setup(self, stage=None):
        self.train_set, self.valid_set = self.ds.split_train_valid()

    def train_dataloader(self):
        return self.train_set.dataloader(
            batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        return self.valid_set.dataloader(
            batch_size=self.batch_size*2, shuffle=False)
data_module = Seq2SeqData(
    df, encoder_tokenizer,
    decoder_tokenizer,
    batch_size=28,
    max_len=256,
    no_punkt=False if TO_CLASSICAL else True,)
data_module.setup()
inputs = next(iter(data_module.train_dataloader()))
inputs<jupyter_output><empty_output><jupyter_text>if we are doing clasical Chinese to modern Chinese, we can randomly set half of the input without any punctuation, as many data source might be<jupyter_code>encoder_tokenizer.batch_decode(
    inputs.input_ids,skip_special_tokens=True
)<jupyter_output><empty_output><jupyter_text>Load pretrained models<jupyter_code># encoder = AutoModel.from_pretrained(ENCODER_PRETRAINED, proxies={"http":"bifrost:3128"})
# decoder = AutoModelForCausalLM.from_pretrained(DECODER_PRETRAINED, add_cross_attention=True,
#                                                proxies={"http":"bifrost:3128"})<jupyter_output><empty_output><jupyter_text>Model We create a seq2seq model by using pretrained encoder + pretrained decoder<jupyter_code># loading pretrained model
encoder_decoder = EncoderDecoderModel.from_encoder_decoder_pretrained(
    encoder_pretrained_model_name_or_path=ENCODER_PRETRAINED,
    decoder_pretrained_model_name_or_path=DECODER_PRETRAINED,
)
class Seq2SeqTrain(pl.LightningModule):
    def __init__(self, encoder_decoder):
        super().__init__()
        self.encoder_decoder = encoder_decoder
        
    def forward(self, batch):
        return self.encoder_decoder(
                **batch
            )

    def training_step(self, batch, batch_idx):
        outputs = self(batch)
        self.log('loss', outputs.loss)
        return outputs.loss

    def validation_step(self, batch, batch_idx):
        outputs = self(batch)
        self.log('val_loss', outputs.loss)
        return outputs.loss
    
    def configure_optimizers(self):
        encoder_params = list(
            {"params":param,"lr":1e-5}
            for param in self.encoder_decoder.encoder.embeddings.parameters()) +\
            list({"params":param,"lr":1e-5}
            for param in self.encoder_decoder.encoder.encoder.parameters()) +\
            list({"params":param,"lr":1e-3}
            for param in self.encoder_decoder.encoder.pooler.parameters())

        decoder_params = list()
        for name, param in self.encoder_decoder.decoder.named_parameters():
            if 'ln_cross_attn' in name:
                decoder_params.append({"params":param,"lr":1e-3})
            elif 'crossattention' in name:
                decoder_params.append({"params":param,"lr":1e-3})
            elif 'lm_head' in name:
                decoder_params.append({"params":param,"lr":1e-4})
            else:
                decoder_params.append({"params":param,"lr":1e-5})

        return torch.optim.Adam(
                encoder_params + decoder_params,
                lr=1e-3,
            )
module = Seq2SeqTrain(encoder_decoder)<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>save = pl.callbacks.ModelCheckpoint(
    '/GCI/transformers/weights/cc_to_zh',
    save_top_k=2,
    verbose=True,
    monitor='val_loss',
    mode='min',
)

trainer = pl.Trainer(
    gpus=[1],
    max_epochs=10,
    callbacks=[save],
)
trainer.fit(module, datamodule=data_module)<jupyter_output>LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name            | Type                | Params
--------------------------------------------------------
0 | encoder_decoder | EncoderDecoderModel | 233 M 
--------------------------------------------------------
233 M     Trainable params
0         Non-trainable params
233 M     Total params
935.203   Total estimated model params size (MB)
{'fraction_non_alphanumeric': 0.110082611724893, 'fraction_numerical': 0.009057430078630437, 'mean_word_length': 2.659140568099053, 'pattern_counts': {'":': 18, '<': 40, '<?xml version=': 0, '>': 41, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '509851', 'license': 'MIT', 'n_tokens_mistral': 3304, 'n_tokens_neox': 3203, 'n_words': 601, 'path': 'nbs/cc2zh_translate.ipynb', 'repo_name': 'raynardj/yuan', 'size': '10047'}