{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4d4f47fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!python3 -m spacy download en_core_web_trf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3d419190",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "import pandas as pd\n",
    "import spacy\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c1e5fde4",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"en_core_web_trf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "98b3d29e",
   "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></th>\n",
       "      <th>Unnamed: 0</th>\n",
       "      <th>s1</th>\n",
       "      <th>s2</th>\n",
       "      <th>label</th>\n",
       "      <th>corrected_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>PCCW's chief operating officer, Mike Butcher, ...</td>\n",
       "      <td>Current Chief Operating Officer Mike Butcher a...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>The world's two largest automakers said their ...</td>\n",
       "      <td>Domestic sales at both GM and No. 2 Ford Motor...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>According to the federal Centers for Disease C...</td>\n",
       "      <td>The Centers for Disease Control and Prevention...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>A tropical storm rapidly developed in the Gulf...</td>\n",
       "      <td>A tropical storm rapidly developed in the Gulf...</td>\n",
       "      <td>0</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>The company didn't detail the costs of the rep...</td>\n",
       "      <td>But company officials expect the costs of the ...</td>\n",
       "      <td>0</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Unnamed: 0                                                 s1  \\\n",
       "0           0  PCCW's chief operating officer, Mike Butcher, ...   \n",
       "1           1  The world's two largest automakers said their ...   \n",
       "2           2  According to the federal Centers for Disease C...   \n",
       "3           3  A tropical storm rapidly developed in the Gulf...   \n",
       "4           4  The company didn't detail the costs of the rep...   \n",
       "\n",
       "                                                  s2  label corrected_label  \n",
       "0  Current Chief Operating Officer Mike Butcher a...      1            None  \n",
       "1  Domestic sales at both GM and No. 2 Ford Motor...      1            None  \n",
       "2  The Centers for Disease Control and Prevention...      1            None  \n",
       "3  A tropical storm rapidly developed in the Gulf...      0            None  \n",
       "4  But company officials expect the costs of the ...      0            None  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split = \"test\"\n",
    "df = pd.read_csv(\"./mrpc/mrpc_\"+split+\".csv\")\n",
    "df[\"corrected_label\"] = df.shape[0]*[None]\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c63bcacb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rchop(s, suffix):\n",
    "    if s.endswith(suffix):\n",
    "        return s[:-len(suffix)]\n",
    "    return s\n",
    "\n",
    "def collate_entities(doc):\n",
    "    output = defaultdict(list)\n",
    "    for ent in doc.ents:\n",
    "        ent_text = ent.text\n",
    "        if ent_text[0]=='\"' or ent_text[0]==\"'\":\n",
    "            ent_text = ent_text[1:]\n",
    "        if ent_text[-1]=='\"' or ent_text[-1]==\"'\":\n",
    "            ent_text = ent_text[:-1]\n",
    "        strip_list = [\"percent\", \"%\"]\n",
    "        for s in strip_list:\n",
    "            ent_text = rchop(ent_text, s)\n",
    "        ent_text = ent_text.strip()\n",
    "        if ent_text not in output[ent.label_]:\n",
    "            output[ent.label_].append(ent_text)\n",
    "    return dict(output)\n",
    "\n",
    "replacements = []\n",
    "black_list = ['US->U.S.', 'U.S.->US',\n",
    "              'U.N.->United Nations', 'United Nations->U.N.',\n",
    "              'each year->annually', 'annually->each year']\n",
    "\n",
    "def correct_row(row, verbose=False, replacement=\"longer\"):\n",
    "    global replacements, black_list\n",
    "    s1_text, s2_text = row[\"s1\"], row[\"s2\"]\n",
    "    s1, s2 = nlp(s1_text), nlp(s2_text)\n",
    "    s1_ents = collate_entities(s1)\n",
    "    s2_ents = collate_entities(s2)\n",
    "    ent_types = list(set(list(s1_ents.keys())+list(s2_ents.keys())))\n",
    "    correctable = True\n",
    "    need_correction = False\n",
    "    sub_source = []\n",
    "    sub_target = []\n",
    "    sub_type = []\n",
    "    for et in ent_types:\n",
    "        try:\n",
    "            s1_count = len(s1_ents[et])\n",
    "        except KeyError:\n",
    "            s1_count = 0\n",
    "        try:\n",
    "            s2_count = len(s2_ents[et])\n",
    "        except KeyError:\n",
    "            s2_count = 0\n",
    "        if s1_count != s2_count:\n",
    "            correctable = False\n",
    "            need_correction = True\n",
    "        elif s1_count==s2_count==1:\n",
    "            e1 = s1_ents[et][0]\n",
    "            e2 = s2_ents[et][0]\n",
    "            if e1.lower()!=e2.lower():\n",
    "                need_correction = True\n",
    "                if len(e1) < len(e2):\n",
    "                    sub_source.append(e1)\n",
    "                    sub_target.append(e2)\n",
    "                else:\n",
    "                    sub_source.append(e2)\n",
    "                    sub_target.append(e1)\n",
    "                sub_type.append(et)\n",
    "        elif s1_count>1 or s2_count>1:\n",
    "            ignore_list = set([])\n",
    "            sub_list = set([])\n",
    "            for ent in set(s1_ents[et]+s2_ents[et]):\n",
    "                if ent in s1_ents[et] and ent in s2_ents[et]:\n",
    "                    ignore_list.add(ent)\n",
    "                else:\n",
    "                    need_correction = True\n",
    "                    sub_list.add(ent)\n",
    "            if len(sub_list)==2:\n",
    "                #print(sub_list)\n",
    "                e1, e2 = sub_list\n",
    "                need_correction = True\n",
    "                if e1.lower()!=e2.lower():\n",
    "                    if len(e1) < len(e2):\n",
    "                        sub_source.append(e1)\n",
    "                        sub_target.append(e2)\n",
    "                    else:\n",
    "                        sub_source.append(e2)\n",
    "                        sub_target.append(e1)\n",
    "                    sub_type.append(et)\n",
    "            elif len(sub_list)==0:\n",
    "                pass\n",
    "            else:\n",
    "                correctable = False\n",
    "    #if len(sub_source)>2:\n",
    "    #    verbose = True\n",
    "    if correctable and len(sub_source)>0:\n",
    "        if verbose:\n",
    "            print(s1_text)\n",
    "            print(s2_text)\n",
    "        for s, t, et in zip(sub_source, sub_target, sub_type):\n",
    "            if verbose:\n",
    "                print(et, \":\", s, \"->\", t)\n",
    "            if replacement==\"longer\":\n",
    "                replacement = s+\"->\"+t\n",
    "                if replacement not in black_list:\n",
    "                    s1_text = s1_text.replace(s, t)\n",
    "                    s2_text = s2_text.replace(s, t)\n",
    "                    replacements.append(replacement)\n",
    "            else:\n",
    "                replacement = t+\"->\"+s\n",
    "                if replacement not in black_list:\n",
    "                    s1_text = s1_text.replace(t, s)\n",
    "                    s2_text = s2_text.replace(t, s)\n",
    "                    replacements.append(t+\"->\"+s)\n",
    "        if verbose:\n",
    "            print(\"\")\n",
    "        corrected = True\n",
    "    else:\n",
    "        corrected = False\n",
    "    if corrected:\n",
    "        assert need_correction\n",
    "    return {\"need_correction\": need_correction,\n",
    "            \"corrected\": corrected,\n",
    "            \"s1\": s1_text,\n",
    "            \"s2\": s2_text}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5f61a97a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1725/1725 [03:03<00:00,  9.40it/s]\n"
     ]
    }
   ],
   "source": [
    "total = 0\n",
    "cant_correct = 0\n",
    "corrected = 0\n",
    "no_need_correct = 0\n",
    "\n",
    "new_s1 = []\n",
    "new_s2 = []\n",
    "new_label = []\n",
    "\n",
    "for index, row in tqdm(df.iterrows(), total=df.shape[0]):\n",
    "    if row[\"label\"] == 1:\n",
    "        total += 1\n",
    "        out = correct_row(row, verbose=False, replacement=\"longer\")\n",
    "        if out[\"need_correction\"] and out[\"corrected\"]:\n",
    "            new_s1.append(out[\"s1\"])\n",
    "            new_s2.append(out[\"s2\"])\n",
    "            new_label.append(1)\n",
    "            corrected += 1\n",
    "            out = correct_row(row, verbose=False, replacement=\"shorter\")\n",
    "            new_s1.append(out[\"s1\"])\n",
    "            new_s2.append(out[\"s2\"])\n",
    "            new_label.append(1)\n",
    "        elif out[\"need_correction\"] and not out[\"corrected\"]:\n",
    "            new_s1.append(row[\"s1\"])\n",
    "            new_s2.append(row[\"s2\"])\n",
    "            new_label.append(0)\n",
    "            cant_correct += 1\n",
    "        else:\n",
    "            new_s1.append(row[\"s1\"])\n",
    "            new_s2.append(row[\"s2\"])\n",
    "            new_label.append(1)\n",
    "            no_need_correct += 1\n",
    "    else:\n",
    "        new_s1.append(row[\"s1\"])\n",
    "        new_s2.append(row[\"s2\"])\n",
    "        new_label.append(0)\n",
    "        no_need_correct += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1472f875",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1725\n",
      "835\n",
      "0.13595505617977527\n",
      "121\n"
     ]
    }
   ],
   "source": [
    "print(total)\n",
    "print(no_need_correct)\n",
    "print(corrected+cant_correct)\n",
    "print(corrected/(corrected+cant_correct))\n",
    "print(corrected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "adfbaa3e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'from collections import Counter\\nreplacements_counter = Counter(replacements)\\nreplacements_counter.most_common(30)'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"from collections import Counter\n",
    "replacements_counter = Counter(replacements)\n",
    "replacements_counter.most_common(30)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d740784e",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_out = pd.DataFrame(list(zip(new_s1, new_s2, new_label)),\n",
    "                    columns =[\"s1\", \"s2\", \"label\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3a287b2d",
   "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></th>\n",
       "      <th>s1</th>\n",
       "      <th>s2</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>PCCW's chief operating officer, Mike Butcher, ...</td>\n",
       "      <td>Current Chief Operating Officer Mike Butcher a...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>The world's two largest automakers said their ...</td>\n",
       "      <td>Domestic sales at both GM and No. 2 Ford Motor...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>According to the federal The Centers for Disea...</td>\n",
       "      <td>The The Centers for Disease Control and Preven...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>According to the federal Centers for Disease C...</td>\n",
       "      <td>Centers for Disease Control and Prevention sai...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>A tropical storm rapidly developed in the Gulf...</td>\n",
       "      <td>A tropical storm rapidly developed in the Gulf...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>The company didn't detail the costs of the rep...</td>\n",
       "      <td>But company officials expect the costs of the ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>The settling companies would also assign their...</td>\n",
       "      <td>Under the agreement, the settling companies wi...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>Air Commodore Quaife said the Hornets remained...</td>\n",
       "      <td>Air Commodore John Quaife said the security op...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>A Washington County man may have the countys f...</td>\n",
       "      <td>The countys first and only human case of West ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>Moseley and a senior aide delivered their summ...</td>\n",
       "      <td>General Moseley and a senior aide presented th...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  s1  \\\n",
       "0  PCCW's chief operating officer, Mike Butcher, ...   \n",
       "1  The world's two largest automakers said their ...   \n",
       "2  According to the federal The Centers for Disea...   \n",
       "3  According to the federal Centers for Disease C...   \n",
       "4  A tropical storm rapidly developed in the Gulf...   \n",
       "5  The company didn't detail the costs of the rep...   \n",
       "6  The settling companies would also assign their...   \n",
       "7  Air Commodore Quaife said the Hornets remained...   \n",
       "8  A Washington County man may have the countys f...   \n",
       "9  Moseley and a senior aide delivered their summ...   \n",
       "\n",
       "                                                  s2  label  \n",
       "0  Current Chief Operating Officer Mike Butcher a...      0  \n",
       "1  Domestic sales at both GM and No. 2 Ford Motor...      0  \n",
       "2  The The Centers for Disease Control and Preven...      1  \n",
       "3  Centers for Disease Control and Prevention sai...      1  \n",
       "4  A tropical storm rapidly developed in the Gulf...      0  \n",
       "5  But company officials expect the costs of the ...      0  \n",
       "6  Under the agreement, the settling companies wi...      1  \n",
       "7  Air Commodore John Quaife said the security op...      0  \n",
       "8  The countys first and only human case of West ...      0  \n",
       "9  General Moseley and a senior aide presented th...      0  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_out.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6dda56a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_out.to_csv(\"./mrpc_\"+split+\"_corrected.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77b33d7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
