lstm-crf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 481,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import os\n",
    "import re\n",
    "from collections import Counter\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence\n",
    "# !pip install git+https://github.com/yumoh/torchcrf.git\n",
    "from torchcrf import CRF\n",
    "from tqdm import tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 482,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_path = './train.txt'\n",
    "# test_data_path = './test.txt'\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 483,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_texts, train_tags = [], []\n",
    "with open(train_data_path, 'rt', encoding='utf-8') as f:\n",
    "    for line in f:\n",
    "        if line is None or len(line) == 0:\n",
    "            continue\n",
    "        splits = line.split('/')\n",
    "        train_texts.append(splits[0].strip())\n",
    "        train_tags.append(splits[1].strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 512,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['江苏南通启东市黄金海滩景区',\n",
       "  '江苏省镇江市句容市边城镇赵庄村150号',\n",
       "  '江苏镇江新区金港大道98号',\n",
       "  '江苏省镇江市句容某某镇莲花新村150号',\n",
       "  '江苏省苏州市常熟市新海路月亮小区17幢1号'],\n",
       " ['aibiciidiiiii',\n",
       "  'aiibiiciidiifiiiiii',\n",
       "  'aibicieiiifii',\n",
       "  'aiibiicidiifiiiiiii',\n",
       "  'aiibiiciieiifiiiiiiii'])"
      ]
     },
     "execution_count": 512,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_texts[:5], train_tags[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 485,
   "metadata": {},
   "outputs": [],
   "source": [
    "counter = Counter(''.join(train_texts))\n",
    "char2id = {c: i for i, (c, _) in enumerate(counter.most_common(None), 1)}  # 0 for pad(mask)\n",
    "# a:province, b:city, c:district, d:town, e:road, f:home\n",
    "tag2id = {'o': 0, 'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5, 'f': 6, 'i': 7}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 486,
   "metadata": {},
   "outputs": [],
   "source": [
    "def text2seq(text, char2id):\n",
    "    return [char2id.get(c, 0) for c in text]\n",
    "\n",
    "def tag2seq(tag, tag2id):\n",
    "    return [tag2id.get(c, 0) for c in tag]\n",
    "\n",
    "def padding(l, pad_id=0):\n",
    "    # 输入：[[1, 1, 1], [2, 2], [3]]\n",
    "    # 返回：[(1, 2, 3), (1, 2, 0), (1, 0, 0)] 返回已经是转置后的 [L, B]\n",
    "    return list(itertools.zip_longest(*l, fillvalue=pad_id))\n",
    "\n",
    "def masking(l, pad_id=0):\n",
    "    # 将targets里非pad部分标记为1，pad部分标记为0\n",
    "    m = []\n",
    "    for i, seq in enumerate(l):\n",
    "        m.append([])\n",
    "        for token in seq:\n",
    "            if token == pad_id:\n",
    "                m[i].append(0)\n",
    "            else:\n",
    "                m[i].append(1)\n",
    "    return m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 487,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_gen(texts, tags, batch_size, char2id, tag2id):\n",
    "    X_Y, X, Y = [], [], []\n",
    "    i = 0\n",
    "    for s, t in zip(texts, tags):\n",
    "        x = text2seq(s, char2id)\n",
    "        y = tag2seq(t, tag2id)\n",
    "        if len(x) > 0 and len(y) > 0 and len(x) == len(y):\n",
    "            i += 1\n",
    "            X_Y.append((x, y))\n",
    "        if len(X_Y) == batch_size or i == len(texts):\n",
    "            X_Y.sort(key=lambda x: len(x[0]), reverse=True)\n",
    "            X, Y = zip(*X_Y)\n",
    "            X = padding(X)\n",
    "            mask = masking(X)\n",
    "            Y = padding(Y)\n",
    "            yield torch.tensor(X, dtype=torch.long).to(device), \\\n",
    "                  torch.tensor(Y, dtype=torch.long).to(device), \\\n",
    "                  torch.tensor(mask, dtype=torch.uint8).to(device)\n",
    "            X_Y = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 488,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = data_gen(train_texts, train_tags, 10, char2id, tag2id)\n",
    "# for X, Y, mask in gen:\n",
    "#     print(X.size(), Y.size(), mask.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 489,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BLSTM_CRF(nn.Module):\n",
    "    def __init__(self, vocab_size, tag_size, embedding_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embedding_size = embedding_size\n",
    "        self.embedding = nn.Embedding(vocab_size + 1, embedding_size)\n",
    "        self.lstm = nn.LSTM(embedding_size, hidden_size, bidirectional=True)\n",
    "        self.fc = nn.Linear(2 * hidden_size, tag_size)\n",
    "        self.crf = CRF(tag_size)\n",
    "\n",
    "    def _get_features(self, x, mask=None):\n",
    "        x = self.embedding(x)  # [L, B, E]\n",
    "        if mask is None:\n",
    "            x, _ = self.lstm(x)  # [L, B, 2H]\n",
    "        else:\n",
    "            lengths = [torch.sum(mask[:, i] > 0).item() for i in range(mask.size(1))]\n",
    "            packed = pack_padded_sequence(x, lengths)\n",
    "            x, _ = self.lstm(packed)\n",
    "            x, _ = pad_packed_sequence(x)\n",
    "        x = F.elu(self.fc(x))  # [L, B, T]\n",
    "        return x\n",
    "\n",
    "    def get_loss(self, x, y, mask=None):\n",
    "        x = self._get_features(x, mask)\n",
    "        loss = self.crf(x, y, mask=mask)\n",
    "        return -loss \n",
    "\n",
    "    def decode(self, x, mask=None):\n",
    "        x = self._get_features(x)  # [L, B, 2H]\n",
    "        x = self.crf.decode(x, mask)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 490,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, optimizer):\n",
    "    for i in range(150):\n",
    "#         print('epoch:', i)\n",
    "        for X, Y, mask in data_gen(train_texts, train_tags, 5, char2id, tag2id):\n",
    "            model.zero_grad()\n",
    "            loss = model.get_loss(X, Y, mask)\n",
    "            print('loss:{:.4f}'.format(loss.item()), end='\\r')\n",
    "            loss.backward()\n",
    "            optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 491,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BLSTM_CRF(len(char2id), len(tag2id), 128, 128).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 492,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:0.033310\r"
     ]
    }
   ],
   "source": [
    "train(model, optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 507,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(text, model):\n",
    "    seq = text2seq(text, char2id)\n",
    "    x = torch.tensor(seq, dtype=torch.long).view(-1, 1).to(device)\n",
    "    tags = model.decode(x)[0]\n",
    "    tags = tags.numpy().tolist()\n",
    "    print(tags)\n",
    "    pieces = {}\n",
    "    i = 0\n",
    "    while i < len(tags):\n",
    "        s = i\n",
    "        if tags[i] == 1:\n",
    "            key = 'provice'\n",
    "        elif tags[i] == 2:\n",
    "            key = 'city'\n",
    "        elif tags[i] == 3:\n",
    "            key = 'country'\n",
    "        elif tags[i] == 4:\n",
    "            key = 'town'\n",
    "        elif tags[i] == 5:\n",
    "            key = 'road'\n",
    "        elif tags[i] == 6:\n",
    "            key = 'house'\n",
    "        else:\n",
    "            i += 1\n",
    "            continue\n",
    "        j = i + 1\n",
    "        while j < len(tags) and tags[j] == 7:\n",
    "            j += 1\n",
    "        pieces[key] = text[s: j]\n",
    "        i = j\n",
    "    return pieces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 508,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3, 7, 7, 5, 7, 7, 6, 7, 7, 7, 7]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'country': '常熟市', 'road': '某某路', 'house': '62号附近'}"
      ]
     },
     "execution_count": 508,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict('常熟市某某路62号附近', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 509,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 7, 0, 0, 0, 0, 2, 7, 0, 0]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'provice': '江苏', 'city': '南通'}"
      ]
     },
     "execution_count": 509,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict('江苏苏州银行南通分行', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 510,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3, 7, 3, 3, 3, 7, 7, 7, 4, 7, 7, 6, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'country': '张家港市', 'town': '锦丰镇', 'house': '锦都花苑15幢103-105室'}"
      ]
     },
     "execution_count": 510,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict('我家住在张家港市锦丰镇锦都花苑15幢103-105室', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 511,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 1, 1, 7, 7, 3, 7, 3, 3, 3, 3, 0, 0, 0, 2, 7, 0, 3, 3, 3, 3, 3, 7, 7, 5, 7, 7, 7, 6, 7, 7, 7]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'provice': '江苏省',\n",
       " 'country': '姑苏区',\n",
       " 'city': '苏州',\n",
       " 'road': '林泉街道',\n",
       " 'house': '189号'}"
      ]
     },
     "execution_count": 511,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict('位于江苏省“人间天堂”之称的苏州，地址为：姑苏区林泉街道189号', model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}