Deep Learning/訓練 LSTM 之歌詞自動填詞器/lstm_writer_hw.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rSmJJczZo5zp"
   },
   "source": [
    "# 專題（一）：訓練LSTM之歌詞自動填詞器\n",
    "\n",
    "## 專案目標\n",
    "- 目標：使用 LSTM 模型去學習五月天歌詞，並且可以自動填詞\n",
    "- mayday_lyrics.txt 資料說明：\n",
    "    - 每一行都是一首歌的歌詞\n",
    "    - 除去標點符號並以空白表示間隔\n",
    "- 利用 mayday_lyrics.txt 來產生歌詞序列\n",
    "- 使用 LSTM 模型去學習歌詞序列\n",
    "- 當給定開頭的一段歌詞，例如：”給我一首歌”，就可以用 LSTM 猜下一個字，反覆這個過程就可以自動填詞\n",
    "\n",
    "## 實作提示\n",
    "- STEP1：從 mayday_lyrics.txt 中取出歌詞\n",
    "- STEP2：建立每個字的 Index\n",
    "- STEP3：用 Rolling 的方式打造 LyricsDataset\n",
    "- STEP4：使用 DataLoader 來包裝 LyricsDataset\n",
    "- STEP5：建立 LSTM 模型： inputs > nn.Embedding > nn.LSTM > nn.Dropout > 取最後一個 state > nn.Linear > softmax\n",
    "- STEP6：開始訓練並調整參數\n",
    "- STEP7：進行 Demo，給定 pre_text ，使用模型迭代來預測下一個字產生歌詞\n",
    "- (進階) STEP8：在 Demo 時可以採用 Softmax 機率來作隨機採樣，這可以增加隨機性，讓歌詞有更多變化，當然還可以使用機率閥值來避免太奇怪的字出現\n",
    "\n",
    "## 重要知識點：專題結束後可以學會\n",
    "- 如何讀取並處理需要 Rolling 的序列資料\n",
    "- 了解如何用 Pytorch 建制一個 LSTM 的模型\n",
    "- 學會如何訓練一個語言模型\n",
    "- 學會如何隨機抽樣 Softmax 的分布"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "hen1MQ1F_cly"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "7joX8NEEu90J"
   },
   "outputs": [],
   "source": [
    "# data from: https://github.com/gaussic/Chinese-Lyric-Corpus\n",
    "with open('mayday_lyrics.txt', encoding='utf-8') as f:\n",
    "    lyrics_list = [line.strip() for line in f.readlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "_P3Bonv_7FpS"
   },
   "outputs": [],
   "source": [
    "# 建立詞典對照表\n",
    "cnt = Counter(''.join(lyrics_list))\n",
    "word2index = {word: idx for idx, word in enumerate(cnt)}\n",
    "index2word = {idx: word for word, idx in word2index.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "nueOEx287Hpm",
    "outputId": "61ce5cc2-7693-4138-8e5c-cc3a56e97ad0"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2101"
      ]
     },
     "execution_count": 4,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(word2index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "BwIpjwU_8YJB"
   },
   "outputs": [],
   "source": [
    "# 建立數據集\n",
    "class LyricsDataset(Dataset):\n",
    "    def __init__(self, lyrics_list, word2index, num_unrollings=10):\n",
    "        self.word2index = word2index\n",
    "        self.samples = []\n",
    "        for lyrics in lyrics_list:\n",
    "            for idx in range(len(lyrics) - num_unrollings + 1):\n",
    "                self.samples.append(lyrics[idx:idx + num_unrollings])\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        sample = self.samples[idx]\n",
    "        input_lyric = torch.LongTensor([self.word2index[w] for w in sample[:-1]])\n",
    "        output_lyric = torch.LongTensor([self.word2index[sample[-1]]])\n",
    "\n",
    "        return input_lyric, output_lyric\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "_L7UokS6_W7O"
   },
   "outputs": [],
   "source": [
    "batch_size = 128\n",
    "\n",
    "dataset = LyricsDataset(lyrics_list, word2index)\n",
    "train_loader = DataLoader(\n",
    "    dataset=dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "_HV3bcPM_l7X"
   },
   "outputs": [],
   "source": [
    "# 建立模型\n",
    "class LSTM_LM(nn.Module):\n",
    "    def __init__(self, vocab_size, n_hidden, num_layers, dropout_ratio):\n",
    "        super(LSTM_LM, self).__init__()\n",
    "        self.embed = nn.Embedding(vocab_size, n_hidden)\n",
    "        self.lstm = nn.LSTM(input_size=n_hidden,\n",
    "                            hidden_size=n_hidden,\n",
    "                            num_layers=num_layers,\n",
    "                            batch_first=True,\n",
    "                            dropout=dropout_ratio)\n",
    "        self.fc = nn.Linear(n_hidden, vocab_size)\n",
    "        self.dropout = nn.Dropout(dropout_ratio)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        embed = self.embed(inputs)  # [batch_size, num_unrollings - 1, n_hidden]\n",
    "        outputs, _ = self.lstm(embed)\n",
    "        outputs = self.dropout(outputs)\n",
    "        output = outputs[:,-1]  # [batch_size, n_hidden]\n",
    "        logits = self.fc(output)\n",
    "\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "-tNgfswV_nc1"
   },
   "outputs": [],
   "source": [
    "def train_batch(model, data, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    inputs, targets = [d.to(device) for d in data]\n",
    "    outputs = model(inputs)\n",
    "    loss = criterion(outputs, targets.view(-1))\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "H_QVIJqHRsnX",
    "outputId": "36782a72-1bbb-4cda-e242-6d3ea54c8fd5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, train_loss: 5.665972850566171\n",
      "epoch 2, train_loss: 5.25386716693558\n",
      "epoch 3, train_loss: 4.980346700753231\n",
      "epoch 4, train_loss: 4.73194461736047\n",
      "epoch 5, train_loss: 4.503753290925858\n",
      "epoch 6, train_loss: 4.296388218150082\n",
      "epoch 7, train_loss: 4.0994421856577885\n",
      "epoch 8, train_loss: 3.9187318455128612\n",
      "epoch 9, train_loss: 3.746693789387104\n",
      "epoch 10, train_loss: 3.5832426848032495\n",
      "Example: 摸不到的顏色 是否+是\n",
      "Example:  只留下結果 時間+世\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 11, train_loss: 3.4247585837473506\n",
      "epoch 12, train_loss: 3.280980289677152\n",
      "epoch 13, train_loss: 3.139299197975509\n",
      "epoch 14, train_loss: 3.006912438536116\n",
      "epoch 15, train_loss: 2.886356460821678\n",
      "epoch 16, train_loss: 2.7607782600314437\n",
      "epoch 17, train_loss: 2.6641395160132473\n",
      "epoch 18, train_loss: 2.55301278853636\n",
      "epoch 19, train_loss: 2.4545082605395705\n",
      "epoch 20, train_loss: 2.372064666350731\n",
      "Example: 摸不到的顏色 是否+叫\n",
      "Example:  只留下結果 時間+無\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 21, train_loss: 2.2735611999401324\n",
      "epoch 22, train_loss: 2.2005940740253442\n",
      "epoch 23, train_loss: 2.1241670016580128\n",
      "epoch 24, train_loss: 2.053541174172093\n",
      "epoch 25, train_loss: 1.9862238547637103\n",
      "epoch 26, train_loss: 1.9201504331169796\n",
      "epoch 27, train_loss: 1.8587811070232825\n",
      "epoch 28, train_loss: 1.8016223666481406\n",
      "epoch 29, train_loss: 1.7444541364610835\n",
      "epoch 30, train_loss: 1.6978119059410473\n",
      "Example: 摸不到的顏色 是否+會\n",
      "Example:  只留下結果 時間+真\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 31, train_loss: 1.642223827766776\n",
      "epoch 32, train_loss: 1.6032566453006323\n",
      "epoch 33, train_loss: 1.5583795946434302\n",
      "epoch 34, train_loss: 1.5169846215559843\n",
      "epoch 35, train_loss: 1.4787571782398992\n",
      "epoch 36, train_loss: 1.442080391796404\n",
      "epoch 37, train_loss: 1.4154223728114934\n",
      "epoch 38, train_loss: 1.373276550667352\n",
      "epoch 39, train_loss: 1.3385484995521244\n",
      "epoch 40, train_loss: 1.3117476935163583\n",
      "Example: 摸不到的顏色 是否+叫\n",
      "Example:  只留下結果 時間+甘\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 41, train_loss: 1.2778594080842642\n",
      "epoch 42, train_loss: 1.2500886305641412\n",
      "epoch 43, train_loss: 1.2238723054480583\n",
      "epoch 44, train_loss: 1.198951853335538\n",
      "epoch 45, train_loss: 1.1729550710451144\n",
      "epoch 46, train_loss: 1.1614444809595932\n",
      "epoch 47, train_loss: 1.1224794107329839\n",
      "epoch 48, train_loss: 1.106105827658995\n",
      "epoch 49, train_loss: 1.0841616753265417\n",
      "epoch 50, train_loss: 1.0662319716225945\n",
      "Example: 摸不到的顏色 是否+叫\n",
      "Example:  只留下結果 時間+等\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 51, train_loss: 1.047770702290797\n",
      "epoch 52, train_loss: 1.0331407420744343\n",
      "epoch 53, train_loss: 1.0150423911055184\n",
      "epoch 54, train_loss: 0.9962996163246647\n",
      "epoch 55, train_loss: 0.9814684118408834\n",
      "epoch 56, train_loss: 0.9639887002539265\n",
      "epoch 57, train_loss: 0.9442754440530274\n",
      "epoch 58, train_loss: 0.938075808513136\n",
      "epoch 59, train_loss: 0.9092470071938362\n",
      "epoch 60, train_loss: 0.9031292626007087\n",
      "Example: 摸不到的顏色 是否+會\n",
      "Example:  只留下結果 時間+結\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 61, train_loss: 0.8854812217388284\n",
      "epoch 62, train_loss: 0.873544145300819\n",
      "epoch 63, train_loss: 0.8593856185951217\n",
      "epoch 64, train_loss: 0.8484407868107078\n",
      "epoch 65, train_loss: 0.8414666249155696\n",
      "epoch 66, train_loss: 0.8238727919214551\n",
      "epoch 67, train_loss: 0.8161872013461458\n",
      "epoch 68, train_loss: 0.8060615818867544\n",
      "epoch 69, train_loss: 0.8030063724379495\n",
      "epoch 70, train_loss: 0.7857866166570137\n",
      "Example: 摸不到的顏色 是否+是\n",
      "Example:  只留下結果 時間+偷\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 71, train_loss: 0.7760122888768154\n",
      "epoch 72, train_loss: 0.7686783145100377\n",
      "epoch 73, train_loss: 0.7582955263375241\n",
      "epoch 74, train_loss: 0.7461540906870896\n",
      "epoch 75, train_loss: 0.7340547248892729\n",
      "epoch 76, train_loss: 0.7330111041838246\n",
      "epoch 77, train_loss: 0.7232278371839278\n",
      "epoch 78, train_loss: 0.7158885688689076\n",
      "epoch 79, train_loss: 0.6993811773937818\n",
      "epoch 80, train_loss: 0.7035749749844176\n",
      "Example: 摸不到的顏色 是否+叫\n",
      "Example:  只留下結果 時間+偷\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 81, train_loss: 0.68553557690172\n",
      "epoch 82, train_loss: 0.6817757044264648\n",
      "epoch 83, train_loss: 0.691167439296645\n",
      "epoch 84, train_loss: 0.6684147139875357\n",
      "epoch 85, train_loss: 0.6644770047762263\n",
      "epoch 86, train_loss: 0.6557509430951862\n",
      "epoch 87, train_loss: 0.6516779911646633\n",
      "epoch 88, train_loss: 0.6451868264583436\n",
      "epoch 89, train_loss: 0.6359767056856415\n",
      "epoch 90, train_loss: 0.6387341541760062\n",
      "Example: 摸不到的顏色 是否+叫\n",
      "Example:  只留下結果 時間+偷\n",
      "Example: 麼多的燦爛的夢 以+為\n",
      "epoch 91, train_loss: 0.6308038399078394\n",
      "epoch 92, train_loss: 0.6235019771025686\n",
      "epoch 93, train_loss: 0.6132863385823623\n",
      "epoch 94, train_loss: 0.6062591842904823\n",
      "epoch 95, train_loss: 0.6057392699962673\n",
      "epoch 96, train_loss: 0.5994504418203199\n",
      "epoch 97, train_loss: 0.5967746212592322\n",
      "epoch 98, train_loss: 0.5980796473604958\n",
      "epoch 99, train_loss: 0.5915583491652447\n",
      "epoch 100, train_loss: 0.581781459513396\n",
      "Example: 摸不到的顏色 是否+看\n",
      "Example:  只留下結果 時間+偷\n",
      "Example: 麼多的燦爛的夢 以+為\n"
     ]
    }
   ],
   "source": [
    "# 訓練模型\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "epochs = 100\n",
    "lr = 0.001\n",
    "\n",
    "model = LSTM_LM(len(word2index), 128, 2, 0.2)\n",
    "model.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(reduction='sum')\n",
    "criterion.to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "for epoch in range(1, 1 + epochs):\n",
    "    tot_train_loss = 0\n",
    "    tot_train_count = 0\n",
    "\n",
    "    for train_data in train_loader:\n",
    "        loss = train_batch(model, train_data, criterion, optimizer, device)\n",
    "        tot_train_loss += loss\n",
    "        tot_train_count += train_data[0].size(0)\n",
    "    print(f\"epoch {epoch}, train_loss: {tot_train_loss / tot_train_count}\")\n",
    "\n",
    "    if epoch % 10 == 0:\n",
    "        for idx in [0, 50, 99]:\n",
    "            input_batch = dataset[idx][0].unsqueeze(0).to(device)\n",
    "            predict = model(input_batch).argmax(dim=-1).item()\n",
    "            print(f\"Example: {dataset.samples[idx][:-1]}+{index2word[predict]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "anMER7TJTWKy",
    "outputId": "13ac821f-b52d-4978-967d-ecb55a19789a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "給我一首歌路上這裡面夢 我想要對對你 知道 突然 氣情的空衣 雨家最後還是我的的永言 愛我 始嘛的手 我的處我\n"
     ]
    }
   ],
   "source": [
    "# 模型inference\n",
    "pre_text = '給我一首歌'\n",
    "generate_len = 50\n",
    "prob_threshold = 0.01\n",
    "\n",
    "result = [word2index[c] for c in pre_text]\n",
    "for _ in range(generate_len):\n",
    "    input_example = torch.LongTensor([result]).to(device)\n",
    "    logit = model(input_example)\n",
    "\n",
    "    prob = F.softmax(logit, dim=-1)\n",
    "    probs = torch.where(prob > prob_threshold, prob, torch.zeros_like(prob))\n",
    "    predict = torch.multinomial(probs, 1).item()\n",
    "    result += [predict]\n",
    "\n",
    "print(''.join([index2word[i] for i in result]))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "lstm_writer_hw.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}