run.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import math\n",
    "import zipfile\n",
    "import numpy as np\n",
    "from collections import Counter, defaultdict\n",
    "from itertools import permutations, chain\n",
    "from gensim.models import Word2Vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_dir = '/workspace/data/tianchi_ruijin/ruijin_round2_train/'\n",
    "test_a_data_dir = '/workspace/data/tianchi_ruijin/ruijin_round2_test_a/'\n",
    "test_b_data_dir = '/workspace/data/tianchi_ruijin/ruijin_round2_test_b/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 一、类与函数定义"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENTITIES = [\n",
    "    \"Amount\", \"Anatomy\", \"Disease\", \"Drug\",\n",
    "    \"Duration\", \"Frequency\", \"Level\", \"Method\",\n",
    "    \"Operation\", \"Reason\", \"SideEff\", \"Symptom\",\n",
    "    \"Test\", \"Test_Value\", \"Treatment\"\n",
    "]\n",
    "\n",
    "RELATIONS = [\n",
    "    \"Test_Disease\",\"Symptom_Disease\",\"Treatment_Disease\",\n",
    "    \"Drug_Disease\",\"Anatomy_Disease\",\"Frequency_Drug\",\n",
    "    \"Duration_Drug\",\"Amount_Drug\",\"Method_Drug\",\"SideEff-Drug\"\n",
    "]\n",
    "\n",
    "class Entity(object):\n",
    "    def __init__(self, ent_id, category, start_pos, end_pos, text):\n",
    "        self.ent_id = ent_id\n",
    "        self.category = category\n",
    "        self.start_pos = start_pos\n",
    "        self.end_pos = end_pos\n",
    "        self.text = text\n",
    "    \n",
    "    def __gt__(self, other):\n",
    "        return self.start_pos > other.start_pos\n",
    "    \n",
    "    def offset(self, offset_val):\n",
    "        return Entity(self.ent_id, \n",
    "                      self.category, \n",
    "                      self.start_pos + offset_val,\n",
    "                      self.end_pos + offset_val,\n",
    "                      self.text)\n",
    "    \n",
    "    def __repr__(self):\n",
    "        fmt = '({ent_id}, {category}, ({start_pos}, {end_pos}), {text})'\n",
    "        return fmt.format(**self.__dict__)\n",
    "\n",
    "class Entities(object):\n",
    "    def __init__(self, ents):\n",
    "        self.ents = sorted(ents)\n",
    "        self.ent_dict = dict(zip([ent.ent_id for ent in ents], ents))\n",
    "    \n",
    "    def __getitem__(self, key):\n",
    "        if isinstance(key, int) or isinstance(key, slice):\n",
    "            return self.ents[key]\n",
    "        else:\n",
    "            return self.ent_dict.get(key, None)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.ents)\n",
    "    \n",
    "    def offset(self, offset_val):\n",
    "        ents = [ent.offset(offset_val) for ent in self.ents]\n",
    "        return Entities(ents)\n",
    "    \n",
    "    def vectorize(self, vec_len, cate2idx):\n",
    "        res_vec = np.zeros(vec_len, dtype=int)\n",
    "        for ent in self.ents:\n",
    "            res_vec[ent.start_pos: ent.end_pos] = cate2idx[ent.category]\n",
    "        return res_vec\n",
    "    \n",
    "    def find_entities(self, start_pos, end_pos):\n",
    "        res = []\n",
    "        for ent in self.ents:\n",
    "            if ent.start_pos > end_pos:\n",
    "                break\n",
    "            sp, ep = (max(start_pos, ent.start_pos), min(end_pos, ent.end_pos))\n",
    "            if ep > sp:\n",
    "                new_ent = Entity(ent.ent_id, ent.category, sp, ep, ent.text[:(ep - sp)])\n",
    "                res.append(new_ent)\n",
    "        return Entities(res)\n",
    "    \n",
    "    def __add__(self, other):\n",
    "        ents = self.ents + other.ents\n",
    "        return Entities(ents)\n",
    "    \n",
    "    def merge(self):\n",
    "        merged_ents = []\n",
    "        for ent in self.ents:\n",
    "            if len(merged_ents) == 0:\n",
    "                merged_ents.append(ent)\n",
    "            elif (merged_ents[-1].end_pos == ent.start_pos and \n",
    "                  merged_ents[-1].category == ent.category):\n",
    "                merged_ent = Entity(ent_id=merged_ents[-1].ent_id, \n",
    "                                    category=ent.category,\n",
    "                                    start_pos=merged_ents[-1].start_pos, \n",
    "                                    end_pos=ent.end_pos, \n",
    "                                    text=merged_ents[-1].text + ent.text)\n",
    "                merged_ents[-1] = merged_ent\n",
    "            else:\n",
    "                merged_ents.append(ent)\n",
    "        return Entities(merged_ents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Relation(object):\n",
    "    def __init__(self, rel_id, category, ent1, ent2):\n",
    "        self.rel_id = rel_id\n",
    "        self.category = category\n",
    "        self.ent1 = ent1\n",
    "        self.ent2 = ent2\n",
    "            \n",
    "    @property\n",
    "    def is_valid(self):\n",
    "        return (isinstance(self.ent1, Entity) and\n",
    "                isinstance(self.ent2, Entity) and\n",
    "                [self.ent1.category, self.ent2.category] == re.split('[-_]', self.category))\n",
    "    \n",
    "    @property\n",
    "    def start_pos(self):\n",
    "        return min(self.ent1.start_pos, self.ent2.start_pos)\n",
    "    \n",
    "    @property\n",
    "    def end_pos(self):\n",
    "        return max(self.ent1.end_pos, self.ent2.end_pos)\n",
    "    \n",
    "    def offset(self, offset_val):\n",
    "        return Relation(self.rel_id, \n",
    "                        self.category, \n",
    "                        self.ent1.offset(offset_val),\n",
    "                        self.ent2.offset(offset_val))\n",
    "    \n",
    "    def __gt__(self, other_rel):\n",
    "        return self.ent1.start_pos > other_rel.ent1.start_pos\n",
    "    \n",
    "    def __repr__(self):\n",
    "        fmt = '({rel_id}, {category} Arg1:{ent1} Arg2:{ent2})'\n",
    "        return fmt.format(**self.__dict__)\n",
    "\n",
    "class Relations(object):\n",
    "    def __init__(self, rels):\n",
    "        self.rels = rels\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        if isinstance(key, int):\n",
    "            return self.rels[key]\n",
    "        elif isinstance(key, slice):\n",
    "            return Relations(self.rels[key])\n",
    "    \n",
    "    def __add__(self, other):\n",
    "        rels = self.rels + other.rels\n",
    "        return Relations(rels)\n",
    "    \n",
    "    def find_relations(self, start_pos, end_pos):\n",
    "        res = []\n",
    "        for rel in self.rels:\n",
    "            if start_pos <= rel.start_pos and end_pos >= rel.end_pos:\n",
    "                res.append(rel)\n",
    "        return Relations(res)\n",
    "    \n",
    "    def offset(self, offset_val): \n",
    "        return Relations([rel.offset(offset_val) for rel in self.rels])\n",
    "    \n",
    "    @property\n",
    "    def start_pos(self):\n",
    "        return min([rel.start_pos for rel in self.rels])\n",
    "\n",
    "    @property\n",
    "    def end_pos(self):\n",
    "        return max([rel.end_pos for rel in self.rels])\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.rels)\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return self.rels.__repr__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextSpan(object):\n",
    "    def __init__(self, text, ents, rels, **kwargs):\n",
    "        self.text = text\n",
    "        self.ents = ents\n",
    "        self.rels = rels\n",
    "    \n",
    "    def __getitem__(self, key):\n",
    "        if isinstance(key, int):\n",
    "            start, stop = key, key + 1\n",
    "        elif isinstance(key, slice):\n",
    "            start = key.start if key.start is not None else 0\n",
    "            stop = key.stop if key.stop is not None else len(self.text)\n",
    "        else:\n",
    "            raise ValueError('parameter should be int or slice')\n",
    "        if start < 0:\n",
    "            start += len(self.text)\n",
    "        if stop < 0:\n",
    "            stop += len(self.text)\n",
    "        text = self.text[key]\n",
    "        ents = self.ents.find_entities(start, stop).offset(-start)\n",
    "        rels = self.rels.find_relations(start, stop).offset(-start)\n",
    "        return TextSpan(text, ents, rels)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.text)\n",
    "    \n",
    "\n",
    "class Sentence(object):\n",
    "    def __init__(self, doc_id, offset, text='', ents=[], rels=[], textspan=None):\n",
    "        self.doc_id = doc_id\n",
    "        self.offset = offset\n",
    "        if isinstance(textspan, TextSpan):\n",
    "            self.textspan = textspan\n",
    "        else:\n",
    "            self.textspan = TextSpan(text, ents, rels)\n",
    "    \n",
    "    @property\n",
    "    def text(self):\n",
    "        return self.textspan.text\n",
    "    \n",
    "    @property\n",
    "    def ents(self):\n",
    "        return self.textspan.ents\n",
    "    \n",
    "    @property\n",
    "    def rels(self):\n",
    "        return self.textspan.rels\n",
    "    \n",
    "    def abbreviate(self, max_len, ellipse_chars='$$'):\n",
    "        if max_len <= len(ellipse_chars):\n",
    "            return ''\n",
    "        left_trim = (max_len - len(ellipse_chars)) // 2\n",
    "        right_trim = max_len - len(ellipse_chars) - left_trim\n",
    "        return self[:left_trim] + ellipse_chars + self[-right_trim:]\n",
    "    \n",
    "    def __getitem__(self, key):\n",
    "        if isinstance(key, int):\n",
    "            start, stop = key, key + 1\n",
    "        elif isinstance(key, slice):\n",
    "            start = key.start if key.start is not None else 0\n",
    "            stop = key.stop if key.stop is not None else len(self.text)\n",
    "        else:\n",
    "            raise ValueError('parameter should be int or slice')\n",
    "        if start < 0:\n",
    "            start += len(self.text)\n",
    "        if stop < 0:\n",
    "            stop += len(self.text)\n",
    "        offset = self.offset + start\n",
    "        textspan = self.textspan[start: stop]\n",
    "        return Sentence(self.doc_id, offset, textspan=textspan)\n",
    "    \n",
    "    def __gt__(self, other):\n",
    "        return self.offset > other.offse\n",
    "    \n",
    "    def __add__(self, other):\n",
    "        if isinstance(other, str):\n",
    "            return Sentence(doc_id=self.doc_id, offset=self.offset, text=self.text + other, \n",
    "                            ents=self.ents, rels=self.rels)\n",
    "        assert self.doc_id == other.doc_id, 'sentences should be from the same document'\n",
    "        assert self.offset + len(self) <= other.offset, 'sentences should not have overlap'\n",
    "        doc_id = self.doc_id\n",
    "        text = self.text + other.text\n",
    "        offset = self.offset\n",
    "        ents = self.ents + other.ents.offset(len(self.text))\n",
    "        rels = self.rels + other.rels.offset(len(self.text))\n",
    "        return Sentence(doc_id=doc_id, offset=offset, text=text, ents=ents, rels=rels)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.textspan)\n",
    "    \n",
    "\n",
    "class Document(object):\n",
    "    def __init__(self, doc_id, text, ents, rels):\n",
    "        self.doc_id = doc_id\n",
    "        self.textspan = TextSpan(text, ents, rels)\n",
    "    \n",
    "    @property\n",
    "    def text(self):\n",
    "        return self.textspan.text\n",
    "    \n",
    "    @property\n",
    "    def ents(self):\n",
    "        return self.textspan.ents\n",
    "    \n",
    "    @property\n",
    "    def rels(self):\n",
    "        return self.textspan.rels\n",
    "    \n",
    "\n",
    "class Documents(object):\n",
    "    def __init__(self, data_dir, doc_ids=None):\n",
    "        self.data_dir = data_dir\n",
    "        self.doc_ids = doc_ids\n",
    "        if self.doc_ids is None:\n",
    "            self.doc_ids = self.scan_doc_ids()\n",
    "    \n",
    "    def scan_doc_ids(self):\n",
    "        doc_ids = [fname.split('.')[0] for fname in os.listdir(self.data_dir)]\n",
    "        doc_ids = [doc_id for doc_id in doc_ids if len(doc_id) > 0]\n",
    "        return np.unique(doc_ids)\n",
    "    \n",
    "    def read_txt_file(self, doc_id):\n",
    "        fname = os.path.join(self.data_dir, doc_id + '.txt')\n",
    "        with open(fname, encoding='utf-8') as f:\n",
    "            text = f.read()\n",
    "        return text\n",
    "    \n",
    "    def parse_entity_line(self, raw_str):\n",
    "        ent_id, label, text = raw_str.strip().split('\\t')\n",
    "        category, pos = label.split(' ', 1)\n",
    "        pos = pos.split(' ')\n",
    "        ent = Entity(ent_id, category, int(pos[0]), int(pos[-1]), text)\n",
    "        return ent\n",
    "    \n",
    "    def parse_relation_line(self, raw_str, ents):\n",
    "        rel_id, label = raw_str.strip().split('\\t')\n",
    "        category, arg1, arg2 = label.split(' ')\n",
    "        arg1 = arg1.split(':')[1]\n",
    "        arg2 = arg2.split(':')[1]\n",
    "        ent1 = ents[arg1]\n",
    "        ent2 = ents[arg2]\n",
    "        return Relation(rel_id, category, ent1, ent2)\n",
    "    \n",
    "    def read_anno_file(self, doc_id):\n",
    "        ents = []\n",
    "        rels = []\n",
    "        fname = os.path.join(self.data_dir, doc_id + '.ann')\n",
    "        with open(fname, encoding='utf-8') as f:\n",
    "            lines = f.readlines()\n",
    "        \n",
    "        for line in lines:\n",
    "            if line.startswith('T'):\n",
    "                ent = self.parse_entity_line(line)\n",
    "                ents.append(ent)\n",
    "        ents = Entities(ents)\n",
    "        \n",
    "        for line in lines:\n",
    "            if line.startswith('R'):\n",
    "                rel = self.parse_relation_line(line, ents)\n",
    "                if rel.is_valid:\n",
    "                    rels.append(rel)\n",
    "        rels = Relations(rels)\n",
    "        return ents, rels\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.doc_ids)\n",
    "    \n",
    "    def get_doc(self, doc_id):\n",
    "        text = self.read_txt_file(doc_id)\n",
    "        ents, rels = self.read_anno_file(doc_id)\n",
    "        doc = Document(doc_id, text, ents, rels)\n",
    "        return doc\n",
    "    \n",
    "    def __getitem__(self, key):\n",
    "        if isinstance(key, int):\n",
    "            doc_id = self.doc_ids[key]\n",
    "            return self.get_doc(doc_id)\n",
    "        if isinstance(key, str):\n",
    "            doc_id = key\n",
    "            return self.get_doc(doc_id)\n",
    "        if isinstance(key, np.ndarray) and key.dtype == int:\n",
    "            doc_ids = self.doc_ids[key]\n",
    "            return Documents(self.data_dir, doc_ids=doc_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SentenceExtractor(object):\n",
    "    def __init__(self, sent_split_char, window_size, rel_types, filter_no_rel_candidates_sents=True):\n",
    "        self.sent_split_char = sent_split_char\n",
    "        self.window_size = window_size\n",
    "        self.filter_no_rel_candidates_sents = filter_no_rel_candidates_sents\n",
    "        self.rels_type_set = set()\n",
    "        for rel_type in rel_types:\n",
    "            self.rels_type_set.add(tuple(re.split('[-_]', rel_type)))\n",
    "        \n",
    "    def get_sent_boundaries(self, text):\n",
    "        dot_indices = []\n",
    "        for i, ch in enumerate(text):\n",
    "            if ch == self.sent_split_char:\n",
    "                dot_indices.append(i + 1)\n",
    "        \n",
    "        if len(dot_indices) <= self.window_size - 1:\n",
    "            return [(0, len(text))]\n",
    "        \n",
    "        dot_indices = [0] + dot_indices\n",
    "        if text[-1] != self.sent_split_char:\n",
    "            dot_indices += [len(text)]\n",
    "\n",
    "        boundries = []\n",
    "        for i in range(len(dot_indices) - self.window_size):\n",
    "            start_stop = (\n",
    "                dot_indices[i],\n",
    "                dot_indices[i + self.window_size]\n",
    "            )\n",
    "            boundries.append(start_stop)\n",
    "        return boundries\n",
    "            \n",
    "    def has_rels_candidates(self, ents):\n",
    "        ent_cates = set([ent.category for ent in ents])\n",
    "        for pos_rel in permutations(ent_cates, 2):\n",
    "            if pos_rel in self.rels_type_set:\n",
    "                return True\n",
    "        return False\n",
    "        \n",
    "    def extract_doc(self, doc):\n",
    "        sents = []\n",
    "        for start_pos, end_pos in self.get_sent_boundaries(doc.text):\n",
    "            ents = []\n",
    "            sent_text = doc.text[start_pos: end_pos]\n",
    "            for ent in doc.ents.find_entities(start_pos=start_pos, end_pos=end_pos):\n",
    "                ents.append(ent.offset(-start_pos))\n",
    "            if self.filter_no_rel_candidates_sents and not self.has_rels_candidates(ents):\n",
    "                continue\n",
    "            rels = []\n",
    "            for rel in doc.rels.find_relations(start_pos=start_pos, end_pos=end_pos):\n",
    "                rels.append(rel.offset(-start_pos))\n",
    "            sent = Sentence(doc.doc_id, \n",
    "                            offset=start_pos, \n",
    "                            text=sent_text, \n",
    "                            ents=Entities(ents),\n",
    "                            rels=Relations(rels))\n",
    "            sents.append(sent)\n",
    "        return sents\n",
    "        \n",
    "    def __call__(self, docs):\n",
    "        sents = []\n",
    "        for doc in docs:\n",
    "            sents += self.extract_doc(doc)\n",
    "        return sents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityPair(object):\n",
    "    def __init__(self, doc_id, sent, from_ent, to_ent):\n",
    "        self.doc_id = doc_id\n",
    "        self.sent = sent\n",
    "        self.from_ent = from_ent\n",
    "        self.to_ent = to_ent\n",
    "        \n",
    "    def __repr__(self):\n",
    "        fmt = 'doc {}, sent {}, {} -> {}'\n",
    "        return fmt.format(self.doc_id, self.sent.text, self.from_ent, self.to_ent)\n",
    "\n",
    "class EntityPairsExtractor(object):\n",
    "    def __init__(self, allow_rel_types, max_len=150, ellipse_chars='$$', pad=10):\n",
    "        self.allow_rel_types = allow_rel_types\n",
    "        self.max_len = max_len\n",
    "        self.pad = pad\n",
    "        self.ellipse_chars = ellipse_chars\n",
    "        \n",
    "    def extract_candidate_rels(self, sent):\n",
    "        candidate_rels = []\n",
    "        for f_ent, t_ent in permutations(sent.ents, 2):\n",
    "            rel_cate = (f_ent.category, t_ent.category)\n",
    "            if rel_cate in self.allow_rel_types:\n",
    "                candidate_rels.append((f_ent, t_ent))\n",
    "        return candidate_rels\n",
    "        \n",
    "    def make_entity_pair(self, sent, f_ent, t_ent):\n",
    "        doc_id = sent.doc_id\n",
    "        if f_ent.start_pos < t_ent.start_pos:\n",
    "            left_ent, right_ent = f_ent, t_ent\n",
    "        else:\n",
    "            left_ent, right_ent = t_ent, f_ent\n",
    "        start_pos = max(0, left_ent.start_pos - self.pad)\n",
    "        end_pos = min(len(sent), right_ent.end_pos + self.pad)\n",
    "        res_sent = sent[start_pos: end_pos]\n",
    "        \n",
    "        if len(res_sent) > self.max_len:\n",
    "            res_sent = res_sent.abbreviate(self.max_len)\n",
    "        f_ent = res_sent.ents[f_ent.ent_id]\n",
    "        t_ent = res_sent.ents[t_ent.ent_id]\n",
    "        return EntityPair(doc_id, res_sent, f_ent, t_ent)\n",
    "    \n",
    "    def __call__(self, sents):\n",
    "        samples = []\n",
    "        for sent in sents:\n",
    "            for f_ent, t_ent in self.extract_candidate_rels(sent.ents):\n",
    "                entity_pair = self.make_entity_pair(sent, f_ent, t_ent)\n",
    "                samples.append(entity_pair)\n",
    "        return samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset(object):\n",
    "    def __init__(self, entity_pairs, doc_ent_pair_ids=set(), word2idx=None, cate2idx=None, max_len=150):\n",
    "        self.entity_pairs = entity_pairs\n",
    "        self.doc_ent_pair_ids = doc_ent_pair_ids\n",
    "        self.max_len = max_len\n",
    "        self.word2idx = word2idx\n",
    "        self.cate2idx = cate2idx\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.entity_pairs)\n",
    "    \n",
    "    def build_vocab_dict(self, vocab_size=2000):\n",
    "        counter = Counter()\n",
    "        for ent_pair in self.entity_pairs:\n",
    "            for char in ent_pair.sent.text:\n",
    "                counter[char] += 1\n",
    "        word2idx = dict()\n",
    "        word2idx['<pad>'] = 0\n",
    "        word2idx['<unk>'] = 1\n",
    "        if vocab_size > 0:\n",
    "            num_most_common = vocab_size - len(word2idx)\n",
    "        else:\n",
    "            num_most_common = len(counter)\n",
    "        for char, _ in counter.most_common(num_most_common):\n",
    "            word2idx[char] = word2idx.get(char, len(word2idx))\n",
    "        self.word2idx = word2idx\n",
    "    \n",
    "    def vectorize(self, ent_pair):\n",
    "        sent_vec = np.zeros(self.max_len, dtype='int')\n",
    "        for i, c in enumerate(ent_pair.sent.text):\n",
    "            sent_vec[i] = self.word2idx.get(c, 1) \n",
    "        ents_vec = ent_pair.sent.ents.vectorize(vec_len=self.max_len, cate2idx=self.cate2idx)\n",
    "        from_ent_vec = np.zeros(self.max_len, dtype='int')\n",
    "        from_ent_vec[ent_pair.from_ent.start_pos: ent_pair.from_ent.end_pos] = 1\n",
    "        to_ent_vec = np.zeros(self.max_len, dtype='int')\n",
    "        to_ent_vec[ent_pair.to_ent.start_pos: ent_pair.to_ent.end_pos] = 1\n",
    "        \n",
    "        if (ent_pair.sent.doc_id, ent_pair.from_ent.ent_id, ent_pair.to_ent.ent_id) in self.doc_ent_pair_ids:\n",
    "            label = 1\n",
    "        else:\n",
    "            label = 0\n",
    "        return sent_vec, ents_vec, from_ent_vec, to_ent_vec, label\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        sent_vecs, ents_vecs, from_ent_vecs, to_ent_vecs, ent_dists, labels = [], [], [], [], [], []\n",
    "        entity_pairs = self.entity_pairs[idx]\n",
    "        if not isinstance(entity_pairs, list):\n",
    "            entity_pairs = [entity_pairs]\n",
    "        for ent_pair in entity_pairs:\n",
    "            sent_vec, ents_vec, from_ent_vec, to_ent_vec, label = self.vectorize(ent_pair)\n",
    "            sent_vecs.append(sent_vec)\n",
    "            ents_vecs.append(ents_vec)\n",
    "            from_ent_vecs.append(from_ent_vec)\n",
    "            to_ent_vecs.append(to_ent_vec)\n",
    "            ent_dists.append(ent_pair.to_ent.start_pos - ent_pair.from_ent.end_pos)\n",
    "            labels.append(label)\n",
    "            \n",
    "        sent_vecs = np.array(sent_vecs)\n",
    "        ents_vecs = np.array(ents_vecs)\n",
    "        from_ent_vecs = np.array(from_ent_vecs)\n",
    "        to_ent_vecs = np.array(to_ent_vecs)\n",
    "        ent_dists = np.array(ent_dists)\n",
    "        labels = np.array(labels)\n",
    "        \n",
    "        return sent_vecs, ents_vecs, from_ent_vecs, to_ent_vecs, ent_dists, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_word_embeddings(entity_pairs, word2idx, *args, **kwargs):\n",
    "    w2v_train_sents = []\n",
    "    for ent_pair in entity_pairs:\n",
    "        w2v_train_sents.append(list(ent_pair.sent.text))\n",
    "    w2v_model = Word2Vec(w2v_train_sents, *args, **kwargs)\n",
    "    word2idx.update({w: i for i, w in enumerate(w2v_model.wv.index2word, start=len(word2idx))})\n",
    "    idx2word = {v: k for k, v in word2idx.items()}\n",
    "    vocab_size = len(word2idx)\n",
    "    w2v_embeddings = np.zeros((len(word2idx), w2v_model.vector_size))\n",
    "    for char, char_idx in word2idx.items():\n",
    "        if char in w2v_model.wv:\n",
    "            w2v_embeddings[char_idx] = w2v_model.wv[char]\n",
    "    return word2idx, idx2word, w2v_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 二、生成训练与测试样本"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_rel_types = set([tuple(re.split('[-_]' ,rel)) for rel in RELATIONS])\n",
    "ent2idx = dict(zip(ENTITIES, range(1, len(ENTITIES) + 1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_docs = Documents(train_data_dir)\n",
    "test_docs = Documents(test_b_data_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. 提取训练集中所有的 relation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_ent_pair_ids = set()\n",
    "for doc in train_docs:\n",
    "    for rel in doc.rels:\n",
    "        doc_ent_pair_id = (doc.doc_id, rel.ent1.ent_id, rel.ent2.ent_id)\n",
    "        doc_ent_pair_ids.add(doc_ent_pair_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. 从文档中抽取句子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "sent_extractor = SentenceExtractor(sent_split_char='。', window_size=2, rel_types=RELATIONS, \n",
    "                                   filter_no_rel_candidates_sents=True)\n",
    "train_sents = sent_extractor(train_docs)\n",
    "test_sents = sent_extractor(test_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'中国成人2型糖尿病HBA1C  c控制目标的专家共识\\n目前,2型糖尿病及其并发症已经成为危害公众\\n健康的主要疾病之一,控制血糖是延缓糖尿病进展及\\n其并发症发生的重要措施之一。虽然HBA1C  。'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_sents[0].text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. 从句子中提取 relation 候选集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_len = 150\n",
    "ent_pair_extractor = EntityPairsExtractor(all_rel_types, max_len=max_len)\n",
    "train_entity_pairs = ent_pair_extractor(train_sents)\n",
    "test_entity_pairs = ent_pair_extractor(test_sents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.利用候选 relation 所在的句子训练字符级别的字向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2idx = {'<pad>': 0, '<unk>': 1}\n",
    "word2idx, idx2word, w2v_embeddings = train_word_embeddings(\n",
    "    entity_pairs=chain(train_entity_pairs, test_entity_pairs),\n",
    "    word2idx=word2idx,\n",
    "    size=100,\n",
    "    iter=10\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.生成训练集和测试集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = Dataset(train_entity_pairs, doc_ent_pair_ids, word2idx=word2idx, max_len=max_len, cate2idx=ent2idx)\n",
    "test_data = Dataset(test_entity_pairs, word2idx=word2idx, max_len=max_len, cate2idx=ent2idx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 三、 建立深度学习模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "from keras import layers\n",
    "from keras import backend as K\n",
    "from keras.layers import Input, Embedding, Lambda\n",
    "from keras.layers import Concatenate, Dense\n",
    "from keras.layers import Conv1D, MaxPool1D, Flatten\n",
    "from keras.models import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_ent_classes = len(ENTITIES) + 1\n",
    "ent_emb_size = 2\n",
    "emb_size = w2v_embeddings.shape[-1]\n",
    "vocab_size = len(word2idx)\n",
    "\n",
    "\n",
    "def build_model():\n",
    "    inp_sent = Input(shape=(max_len,), dtype='int32')\n",
    "    inp_ent = Input(shape=(max_len,), dtype='int32')\n",
    "    inp_f_ent = Input(shape=(max_len,), dtype='float32')\n",
    "    inp_t_ent = Input(shape=(max_len,), dtype='float32')\n",
    "    inp_ent_dist = Input(shape=(1,), dtype='float32')\n",
    "    f_ent = Lambda(lambda x: K.expand_dims(x))(inp_f_ent)\n",
    "    t_ent = Lambda(lambda x: K.expand_dims(x))(inp_t_ent)\n",
    "    \n",
    "    ent_embed = Embedding(num_ent_classes, ent_emb_size)(inp_ent)\n",
    "    sent_embed = Embedding(vocab_size, emb_size, weights=[w2v_embeddings], trainable=False)(inp_sent)\n",
    "    \n",
    "    x = Concatenate()([sent_embed, ent_embed])\n",
    "    x = Conv1D(64, 1, padding='same', activation='relu')(x)\n",
    "    \n",
    "    f_res = layers.multiply([f_ent, x])\n",
    "    t_res = layers.multiply([t_ent, x])\n",
    "    \n",
    "    conv = Conv1D(64, 3, padding='same', activation='relu')\n",
    "    f_x = conv(x)\n",
    "    t_x = conv(x)\n",
    "    f_x = layers.add([f_x, f_res])\n",
    "    t_x = layers.add([t_x, t_res])\n",
    "    \n",
    "    f_res = layers.multiply([f_ent, f_x])\n",
    "    t_res = layers.multiply([t_ent, t_x])\n",
    "    conv = Conv1D(64, 3, padding='same', activation='relu')\n",
    "    f_x = conv(x)\n",
    "    t_x = conv(x)\n",
    "    f_x = layers.add([f_x, f_res])\n",
    "    t_x = layers.add([t_x, t_res])\n",
    "\n",
    "    f_res = layers.multiply([f_ent, f_x])\n",
    "    t_res = layers.multiply([t_ent, t_x])\n",
    "    conv = Conv1D(64, 3, padding='same', activation='relu')\n",
    "    f_x = conv(x)\n",
    "    t_x = conv(x)\n",
    "    f_x = layers.add([f_x, f_res])\n",
    "    t_x = layers.add([t_x, t_res])\n",
    "    \n",
    "    conv = Conv1D(64, 3, activation='relu')\n",
    "    f_x = MaxPool1D(3)(conv(f_x))\n",
    "    t_x = MaxPool1D(3)(conv(t_x))\n",
    "    \n",
    "    conv = Conv1D(64, 3, activation='relu')\n",
    "    f_x = MaxPool1D(3)(conv(f_x))\n",
    "    t_x = MaxPool1D(3)(conv(t_x))\n",
    "    \n",
    "    f_x = Flatten()(f_x)\n",
    "    t_x = Flatten()(t_x)\n",
    "    \n",
    "    x = Concatenate()([f_x, t_x, inp_ent_dist])\n",
    "    x = Dense(256, activation='relu')(x)\n",
    "    x = Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "    model = Model([inp_sent, inp_ent, inp_f_ent, inp_t_ent, inp_ent_dist], x)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_sent, tr_ent, tr_f_ent, tr_t_ent, tr_ent_dist, tr_y = train_data[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "K.clear_session()\n",
    "model = build_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile('adam', loss='binary_crossentropy', metrics=['acc'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "553074/553074 [==============================] - 94s 170us/step - loss: 0.2427 - acc: 0.8959\n",
      "Epoch 2/2\n",
      "553074/553074 [==============================] - 89s 161us/step - loss: 0.1874 - acc: 0.9211\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fc29b6307b8>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(x=[tr_sent, tr_ent, tr_f_ent, tr_t_ent, tr_ent_dist], \n",
    "          y=tr_y, batch_size=64, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.模型预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "63173/63173 [==============================] - 4s 69us/step\n"
     ]
    }
   ],
   "source": [
    "te_sent, te_ent, te_f_ent, te_t_ent, te_ent_dist, te_y = test_data[:]\n",
    "preds = model.predict(x=[te_sent, te_ent, te_f_ent, te_t_ent, te_ent_dist], verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 四、结果输出"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_submission(preds, entity_pairs, threshold):\n",
    "    doc_rels = defaultdict(set)\n",
    "    for p, ent_pair in zip(preds, entity_pairs):\n",
    "        if p >= threshold:\n",
    "            doc_id = ent_pair.doc_id\n",
    "            f_ent_id = ent_pair.from_ent.ent_id\n",
    "            t_ent_id = ent_pair.to_ent.ent_id\n",
    "            category = ent_pair.from_ent.category + '_' + ent_pair.to_ent.category\n",
    "            category = category.replace('SideEff_Drug', 'SideEff-Drug')\n",
    "            doc_rels[doc_id].add((f_ent_id, t_ent_id, category))\n",
    "    submits = dict()\n",
    "    tot_num_rels = 0\n",
    "    for doc_id, rels in doc_rels.items():\n",
    "        output_str = ''\n",
    "        for i, rel in enumerate(rels):\n",
    "            tot_num_rels += 1\n",
    "            line = 'R{}\\t{} Arg1:{} Arg2:{}\\n'.format(i + 1, rel[2], rel[0], rel[1])\n",
    "            output_str += line\n",
    "        submits[doc_id] = output_str\n",
    "    print('Total number of relations: {}. In average {} relations per doc.'.format(tot_num_rels, tot_num_rels / len(submits)))\n",
    "    return submits\n",
    "\n",
    "def output_submission(submit_file, submits, test_dir):\n",
    "    dir_name = os.path.basename(submit_file)\n",
    "    dir_name, _ = os.path.splitext(dir_name)\n",
    "    with zipfile.ZipFile(submit_file, 'w') as zf:\n",
    "        for doc_id, rels_str in submits.items():\n",
    "            fname = '{}.ann'.format(doc_id)\n",
    "            test_file = os.path.join(test_dir, fname)\n",
    "            content = open(test_file, encoding='utf-8').read()\n",
    "            content += rels_str\n",
    "            zf.writestr(os.path.join(dir_name, fname), content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of relations: 9185. In average 170.09259259259258 relations per doc.\n"
     ]
    }
   ],
   "source": [
    "submits = generate_submission(preds, test_entity_pairs, 0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_file = 'submit_181205.zip'\n",
    "output_submission(submit_file, submits, test_b_data_dir)"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}