bea_code/run_bea_conll.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.special import digamma\n",
    "from sklearn.preprocessing import LabelEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_le = LabelEncoder().fit(['B-LOC', 'B-MISC', 'B-ORG', 'B-PER', 'I-LOC', 'I-MISC', 'I-ORG', 'I-PER', 'O'])\n",
    "tag_le = LabelEncoder().fit(['LOC', 'MISC', 'ORG', 'PER', 'O'])\n",
    "\n",
    "num_classes = len(label_le.classes_)\n",
    "num_tags = len(tag_le.classes_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def get_entities(labels):\n",
    "    entities = []\n",
    "    pre_label = 'O'\n",
    "    pre_tag = 'O'\n",
    "    pre_sep = 0\n",
    "    for cur_idx, cur_label in enumerate(np.append(labels, 'O')):\n",
    "        cur_tag = cur_label.split('-')[-1]\n",
    "\n",
    "        if cur_tag != pre_tag or cur_label.startswith('B-'):\n",
    "            if pre_tag != 'O':\n",
    "                entities.append(((pre_sep, cur_idx), pre_tag))\n",
    "            pre_sep = cur_idx\n",
    "\n",
    "        pre_label = cur_label\n",
    "        pre_tag = cur_tag\n",
    "    return entities\n",
    "\n",
    "def get_f1(s1, s2):\n",
    "    return 2*len(s1 & s2) / (len(s1) + len(s2)) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mv_infer(values):\n",
    "    num_classes = values.max() + 1\n",
    "    num_items, num_workers = values.shape\n",
    "    \n",
    "    all_items = np.arange(num_items)\n",
    "    z_ik = np.zeros((num_items, num_classes))\n",
    "\n",
    "    for j in range(num_workers):\n",
    "        z_ik[all_items, values[:, j]] += 1\n",
    "\n",
    "    return z_ik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bea_infer(values, alpha=1, beta_kl=1, prior=True):\n",
    "    num_classes = values.max() + 1\n",
    "    num_items, num_workers = values.shape\n",
    "    \n",
    "    beta_kl = beta_kl * np.ones((num_classes, num_classes))\n",
    "    \n",
    "    z_ik = mv_infer(values)\n",
    "    n_jkl = np.empty((num_workers, num_classes, num_classes))\n",
    "    \n",
    "    last_z_ik = z_ik.copy()\n",
    "    for iteration in range(500):\n",
    "        Eq_log_pi_k = digamma(z_ik.sum(axis=0) + alpha) - digamma(num_items + num_classes*alpha)\n",
    "        \n",
    "        n_jkl[:] = beta_kl\n",
    "        for j in range(num_workers):\n",
    "            for k in range(num_classes):\n",
    "                n_jkl[j, k, :] += np.bincount(values[:, j], z_ik[:, k], minlength=num_classes)        \n",
    "        Eq_log_v_jkl = digamma(n_jkl) - digamma(n_jkl.sum(axis=-1, keepdims=True))\n",
    "        \n",
    "        if prior:\n",
    "            z_ik[:] = Eq_log_pi_k\n",
    "        else:\n",
    "            z_ik.fill(0)\n",
    "\n",
    "        for j in range(num_workers):\n",
    "            z_ik += Eq_log_v_jkl[j, :, values[:, j]]\n",
    "        z_ik -= z_ik.max(axis=-1, keepdims=True)\n",
    "        z_ik = np.exp(z_ik)\n",
    "        z_ik /= z_ik.sum(axis=-1, keepdims=True)\n",
    "        \n",
    "        if np.allclose(z_ik, last_z_ik, atol=1e-3):\n",
    "            break\n",
    "            \n",
    "        last_z_ik[:] = z_ik\n",
    "    return z_ik, iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "a_v, b_v = 1, 1\n",
    "beta_kl = np.eye(num_classes) * (a_v-b_v) + b_v\n",
    "beta_kl_tag = np.eye(num_tags) * (a_v-b_v) + b_v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = './data_conll/'\n",
    "languages = ['de', 'en', 'es', 'nl']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "records = []\n",
    "for language in languages:\n",
    "    dataset = language + '_test'\n",
    "    df_label = pd.read_csv(data_path + dataset + '/label.csv')\n",
    "    df_truth = pd.read_csv(data_path + dataset + '/truth.csv')\n",
    "    \n",
    "    true_entities = set(get_entities(df_truth.truth.values))\n",
    "    for source in df_label.columns:\n",
    "        pred_entities = set(get_entities(label_le.inverse_transform(df_label[source].values)))\n",
    "        f1 = get_f1(true_entities, pred_entities)\n",
    "        records.append((dataset, source, f1))\n",
    "        \n",
    "    z_ik = mv_infer(df_label.values)\n",
    "    pred_entities = set(get_entities(label_le.inverse_transform(z_ik.argmax(axis=-1))))\n",
    "    f1 = get_f1(true_entities, pred_entities)\n",
    "    records.append((dataset, 'MV-tok', f1))\n",
    "    \n",
    "    z_ik, iteration = bea_infer(df_label.values, beta_kl=beta_kl, prior=True)\n",
    "    pred_entities = set(get_entities(label_le.inverse_transform(z_ik.argmax(axis=-1))))\n",
    "    f1 = get_f1(true_entities, pred_entities)\n",
    "    records.append((dataset, 'BEA-tok', f1))\n",
    "    \n",
    "    df_range = pd.DataFrame({source: dict(get_entities(label_le.inverse_transform(df_label[source].values))) \n",
    "                             for source in df_label.columns}).fillna('O')\n",
    "    values_range = np.column_stack([tag_le.transform(df_range[source]) for source in df_range.columns])\n",
    "    \n",
    "    z_ik = mv_infer(values_range)\n",
    "    pred_entities = set([(rng, tag) for (rng, tag) \n",
    "                         in zip(df_range.index.values, tag_le.inverse_transform(z_ik.argmax(axis=-1))) if tag != 'O'])\n",
    "    f1 = get_f1(true_entities, pred_entities)\n",
    "    records.append((dataset, 'MV-ent', f1))\n",
    "    \n",
    "    z_ik, iteration = bea_infer(values_range, beta_kl=beta_kl_tag, prior=True)\n",
    "    pred_entities = set([(rng, tag) for (rng, tag) \n",
    "                         in zip(df_range.index.values, tag_le.inverse_transform(z_ik.argmax(axis=-1))) if tag != 'O'])\n",
    "    f1 = get_f1(true_entities, pred_entities)\n",
    "    records.append((dataset, 'BEA-ent', f1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "df_res = pd.DataFrame.from_records(records, columns=['dataset', 'method', 'f1'])\n",
    "df_pivot = df_res.pivot(index='method', columns='dataset', values='f1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th>dataset</th>\n",
       "      <th>de_test</th>\n",
       "      <th>en_test</th>\n",
       "      <th>es_test</th>\n",
       "      <th>nl_test</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>method</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>BEA-ent</th>\n",
       "      <td>57.7</td>\n",
       "      <td>64.8</td>\n",
       "      <td>63.4</td>\n",
       "      <td>70.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>BEA-tok</th>\n",
       "      <td>58.2</td>\n",
       "      <td>61.2</td>\n",
       "      <td>64.7</td>\n",
       "      <td>70.1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>MV-ent</th>\n",
       "      <td>57.7</td>\n",
       "      <td>64.6</td>\n",
       "      <td>69.0</td>\n",
       "      <td>70.3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>MV-tok</th>\n",
       "      <td>57.4</td>\n",
       "      <td>62.1</td>\n",
       "      <td>66.4</td>\n",
       "      <td>71.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>de</th>\n",
       "      <td>NaN</td>\n",
       "      <td>63.6</td>\n",
       "      <td>55.9</td>\n",
       "      <td>63.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>en</th>\n",
       "      <td>55.4</td>\n",
       "      <td>NaN</td>\n",
       "      <td>65.2</td>\n",
       "      <td>68.6</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>es</th>\n",
       "      <td>46.2</td>\n",
       "      <td>50.3</td>\n",
       "      <td>NaN</td>\n",
       "      <td>58.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>nl</th>\n",
       "      <td>58.4</td>\n",
       "      <td>60.7</td>\n",
       "      <td>65.2</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "dataset  de_test  en_test  es_test  nl_test\n",
       "method                                     \n",
       "BEA-ent     57.7     64.8     63.4     70.3\n",
       "BEA-tok     58.2     61.2     64.7     70.1\n",
       "MV-ent      57.7     64.6     69.0     70.3\n",
       "MV-tok      57.4     62.1     66.4     71.0\n",
       "de           NaN     63.6     55.9     63.5\n",
       "en          55.4      NaN     65.2     68.6\n",
       "es          46.2     50.3      NaN     58.0\n",
       "nl          58.4     60.7     65.2      NaN"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_pivot.round(1)"
   ]
  }
 ],
 "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}