nbs/zh2cc_translate.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Translate model\n",
    "\n",
    "We are using [this nice dataset](https://github.com/BangBOOM/Classical-Chinese)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from forgebox.imports import *\n",
    "from datasets import load_dataset\n",
    "# from fastai.text.all import *\n",
    "from unpackai.nlp import *\n",
    "from tqdm.notebook import tqdm\n",
    "import pytorch_lightning as pl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data=Path(\"/some_location/data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA = Path(data/\"nlp\"/\"zh\"/\"cc_vs_zh\")\n",
    "TO_CLASSICAL = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data\n",
    "\n",
    "### Combine data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_file = list(DATA.rglob(\"data/*\"))\n",
    "\n",
    "\n",
    "def open_file_to_lines(file):\n",
    "    with open(file) as f:\n",
    "        lines = f.read().splitlines()\n",
    "    return lines\n",
    "\n",
    "def pairing_the_file(files,kw):\n",
    "    pairs = []\n",
    "    for file in files:\n",
    "        if kw not in file.name:\n",
    "            file1 = file\n",
    "            file2 = f\"{file}{kw}\"\n",
    "            pairs.append((file1,file2))\n",
    "    return pairs\n",
    "\n",
    "pairs = pairing_the_file(all_file,\"翻译\")\n",
    "\n",
    "def open_pairs(pairs):\n",
    "    chunks = []\n",
    "    for pair in tqdm(pairs, leave=False):\n",
    "        file1,file2 = pair\n",
    "        lines1 = open_file_to_lines(file1)\n",
    "        lines2 = open_file_to_lines(file2)\n",
    "        chunks.append(pd.DataFrame({\"classical\":lines1,\"modern\":lines2}))\n",
    "    return pd.concat(chunks).sample(frac=1.).reset_index(drop=True)\n",
    "\n",
    "data_df = open_pairs(pairs)\n",
    "\n",
    "df = data_df.rename(\n",
    "    columns = dict(\n",
    "        zip([\"modern\",\"classical\"],\n",
    "             [\"source\",\"target\"] if TO_CLASSICAL else [\"target\",\"source\",]))\n",
    ")\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForCausalLM,\n",
    "    AutoModel,\n",
    "    EncoderDecoderModel\n",
    "    )\n",
    "\n",
    "# we find a English parsing encoder, as a pretrained bert is good at understanding english\n",
    "# BERT is short for Bidirectional **Encoder** Representations from Transformers, which consists fully of encoder blocks\n",
    "ENCODER_PRETRAINED = \"bert-base-chinese\"\n",
    "# we find a Chinese writing model for decoder, as decoder is the part of the model that can write stuff\n",
    "DECODER_PRETRAINED = \"uer/gpt2-chinese-poem\"\n",
    "\n",
    "encoder_tokenizer = AutoTokenizer.from_pretrained(ENCODER_PRETRAINED)\n",
    "\n",
    "decoder_tokenizer = AutoTokenizer.from_pretrained(\n",
    "    ENCODER_PRETRAINED # notice we use the BERT's tokenizer here\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytoch Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(Dataset):\n",
    "    def __init__(self, df, tokenizer, target_tokenizer, max_len=128):\n",
    "        super().__init__()\n",
    "        self.df = df\n",
    "        self.tokenizer = tokenizer\n",
    "        self.target_tokenizer = target_tokenizer\n",
    "        self.max_len = max_len\n",
    "        \n",
    "    def __len__(self, ):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return dict(self.df.iloc[idx])\n",
    "\n",
    "    def collate(self, batch):\n",
    "        batch_df = pd.DataFrame(list(batch))\n",
    "        x, y = batch_df.source, batch_df.target\n",
    "        x_batch = self.tokenizer(\n",
    "            list(x),\n",
    "            max_length=self.max_len,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            return_tensors='pt',\n",
    "        )\n",
    "        y_batch = self.target_tokenizer(\n",
    "            list(y),\n",
    "            max_length=self.max_len,\n",
    "            padding='max_length',\n",
    "            truncation=True,\n",
    "            return_tensors='pt',\n",
    "        )\n",
    "        x_batch['decoder_input_ids'] = y_batch['input_ids']\n",
    "        x_batch['labels'] = y_batch['input_ids'].clone()\n",
    "        x_batch['labels'][x_batch['labels'] == self.tokenizer.pad_token_id] = -100\n",
    "        return x_batch\n",
    "\n",
    "    def dataloader(self, batch_size, shuffle=True):\n",
    "        return DataLoader(\n",
    "            self,\n",
    "            batch_size=batch_size,\n",
    "            shuffle=shuffle,\n",
    "            collate_fn=self.collate,\n",
    "        )\n",
    "\n",
    "    def split_train_valid(self, valid_size=0.1):\n",
    "        split_index = int(len(self) * (1 - valid_size))\n",
    "        cls = type(self)\n",
    "        shuffled = self.df.sample(frac=1).reset_index(drop=True)\n",
    "        train_set = cls(\n",
    "            shuffled.iloc[:split_index],\n",
    "            tokenizer=self.tokenizer,\n",
    "            target_tokenizer=self.target_tokenizer,\n",
    "            max_len=self.max_len,\n",
    "        )\n",
    "        valid_set = cls(\n",
    "            shuffled.iloc[split_index:],\n",
    "            tokenizer=self.tokenizer,\n",
    "            target_tokenizer=self.target_tokenizer,\n",
    "            max_len=self.max_len,\n",
    "        )\n",
    "        return train_set, valid_set"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PL datamodule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqData(pl.LightningDataModule):\n",
    "    def __init__(self, df, tokenizer, target_tokenizer, batch_size=12, max_len=128):\n",
    "        super().__init__()\n",
    "        self.df = df\n",
    "        self.ds = Seq2Seq(df, tokenizer, target_tokenizer,max_len=max_len)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.target_tokenizer = target_tokenizer\n",
    "        self.max_len = max_len\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        self.train_set, self.valid_set = self.ds.split_train_valid()\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.train_set.dataloader(\n",
    "            batch_size=self.batch_size, shuffle=True)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.valid_set.dataloader(\n",
    "            batch_size=self.batch_size*2, shuffle=False)\n",
    "\n",
    "data_module = Seq2SeqData(df, encoder_tokenizer, decoder_tokenizer, batch_size=64, )\n",
    "data_module.setup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "next(iter(data_module.train_dataloader()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load pretrained models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# loading pretrained model\n",
    "encoder_decoder = EncoderDecoderModel.from_encoder_decoder_pretrained(\n",
    "    encoder_pretrained_model_name_or_path=ENCODER_PRETRAINED,\n",
    "    decoder_pretrained_model_name_or_path=DECODER_PRETRAINED,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqTrain(pl.LightningModule):\n",
    "    def __init__(self, encoder_decoder):\n",
    "        super().__init__()\n",
    "        self.encoder_decoder = encoder_decoder\n",
    "        \n",
    "    def forward(self, batch):\n",
    "        return self.encoder_decoder(\n",
    "                **batch\n",
    "            )\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        outputs = self(batch)\n",
    "        self.log('loss', outputs.loss)\n",
    "        return outputs.loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        outputs = self(batch)\n",
    "        self.log('val_loss', outputs.loss)\n",
    "        return outputs.loss\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        encoder_params = list(\n",
    "            {\"params\":param,\"lr\":1e-5}\n",
    "            for param in self.encoder_decoder.encoder.embeddings.parameters()) +\\\n",
    "            list({\"params\":param,\"lr\":1e-5}\n",
    "            for param in self.encoder_decoder.encoder.encoder.parameters()) +\\\n",
    "            list({\"params\":param,\"lr\":1e-3}\n",
    "            for param in self.encoder_decoder.encoder.pooler.parameters())\n",
    "\n",
    "        decoder_params = list()\n",
    "        for name, param in self.encoder_decoder.decoder.named_parameters():\n",
    "            if 'ln_cross_attn' in name:\n",
    "                decoder_params.append({\"params\":param,\"lr\":1e-3})\n",
    "            elif 'crossattention' in name:\n",
    "                decoder_params.append({\"params\":param,\"lr\":1e-3})\n",
    "            elif 'lm_head' in name:\n",
    "                decoder_params.append({\"params\":param,\"lr\":1e-4})\n",
    "            else:\n",
    "                decoder_params.append({\"params\":param,\"lr\":1e-5})\n",
    "\n",
    "        return torch.optim.Adam(\n",
    "                encoder_params + decoder_params,\n",
    "                lr=1e-3,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "module = Seq2SeqTrain(encoder_decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "save = pl.callbacks.ModelCheckpoint(\n",
    "    data/'../weights/cc_to_zh',\n",
    "    save_top_k=2,\n",
    "    verbose=True,\n",
    "    monitor='val_loss',\n",
    "    mode='min',\n",
    ")\n",
    "\n",
    "trainer = pl.Trainer(\n",
    "    gpus=[0],\n",
    "    max_epochs=10,\n",
    "    callbacks=[save],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.fit(module, datamodule=data_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = save.best\n",
    "module.load_state_dict(torch.load(best, map_location=\"cpu\")['state_dict'])\n",
    "\n",
    "\n",
    "encoder_decoder = encoder_decoder.cpu()\n",
    "encoder_decoder = encoder_decoder.eval()\n",
    "\n",
    "def inference(text, starter=''):\n",
    "    tk_kwargs = dict(truncation=True, max_length=128, padding=\"max_length\",\n",
    "                     return_tensors='pt')\n",
    "    inputs = encoder_tokenizer([text,],**tk_kwargs)\n",
    "    with torch.no_grad():\n",
    "        return decoder_tokenizer.batch_decode(\n",
    "            encoder_decoder.generate(\n",
    "            inputs.input_ids,\n",
    "            attention_mask=inputs.attention_mask,\n",
    "            num_beams=3,\n",
    "            bos_token_id=101,\n",
    "        ),\n",
    "                                              skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inference('我来跟大家说一句话')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inference(\"这个翻译不是很聪明，因为训练数据不够\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_decoder.push_to_hub(\"raynardj/wenyanwen-chinese-translate-to-ancient\")\n",
    "encoder_tokenizer.push_to_hub(\"raynardj/wenyanwen-chinese-translate-to-ancient\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}