Updata: 단어 사전 저장

모델 로드할 때 단어 사전도 같이 로드 할 수 있도록 모델 학습 후에 사전 저장하는 코드 추가
tukcomCD2024 · Jun 19, 2024 · 1155ce0 · 1155ce0
1 parent dea6fb4
commit 1155ce0
Showing 1 changed file with 29 additions and 1 deletion.
diff --git a/AI/LSTM_attention_test.ipynb b/AI/LSTM_attention_test.ipynb
@@ -37,6 +37,7 @@
         "import os\n",
         "import json\n",
         "import csv\n",
+        "import json\n",
         "\n",
         "TL_sentence_path = '/content/drive/MyDrive/LSTM+attention/sentence_dataTL.csv'\n",
         "VL_sentence_path = '/content/drive/MyDrive/LSTM+attention/sentence_dataVL.csv'\n",
@@ -1207,6 +1208,7 @@
         "        self.name = name\n",
         "        self.word2index = {\"UNK\": 2}\n",
         "        self.word2count = {}\n",
+        "\n",
         "        self.index2word = {0: \"SOS\", 1: \"EOS\", 2: \"UNK\", 3: \"PAD\"}\n",
         "        self.n_words = 4  # SOS, EOS, UNK, PAD\n",
         "\n",
@@ -1224,7 +1226,22 @@
         "            self.word2count[word] += 1\n",
         "\n",
         "    def getWordIndex(self, word):\n",
-        "        return self.word2index.get(word, self.word2index[\"UNK\"])"
+        "        return self.word2index.get(word, self.word2index[\"UNK\"])\n",
+        "\n",
+        "    def saveLang(self, filename):\n",
+        "        with open(filename, 'w', encoding='utf-8') as f:\n",
+        "            json.dump({\n",
+        "                'word2index': self.word2index,\n",
+        "                'index2word': self.index2word,\n",
+        "                'n_words': self.n_words\n",
+        "            }, f, ensure_ascii=False, indent=4)\n",
+        "\n",
+        "    def loadLang(self, filename):\n",
+        "        with open(filename, 'r', encoding='utf-8') as f:\n",
+        "            data = json.load(f)\n",
+        "            self.word2index = data['word2index']\n",
+        "            self.index2word = data['index2word']\n",
+        "            self.n_words = data['n_words']"
       ]
     },
     {
@@ -1587,6 +1604,17 @@
         "trainIters(encoder, decoder, 1000, print_every=100, plot_every=50)  # 적은 수의 iteration으로 실행"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# 모델 학습 후 사전 저장\n",
+        "dialect_lang.saveLang('dialect_lang.json')\n",
+        "standard_lang.saveLang('standard_lang.json')"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,