{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab70c295",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!python3 -m spacy download en_core_web_trf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a82aae05",
   "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": "3483bcba",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"en_core_web_trf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f5cf5d52",
   "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>Amrozi accused his brother, whom he called \"th...</td>\n",
       "      <td>Referring to him as only \"the witness\", Amrozi...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>Yucaipa owned Dominick's before selling the ch...</td>\n",
       "      <td>Yucaipa bought Dominick's in 1995 for $693 mil...</td>\n",
       "      <td>0</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>They had published an advertisement on the Int...</td>\n",
       "      <td>On June 10, the ship's owners had published an...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>Around 0335 GMT, Tab shares were up 19 cents, ...</td>\n",
       "      <td>Tab shares jumped 20 cents, or 4.6%, to set a ...</td>\n",
       "      <td>0</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>The stock rose $2.11, or about 11 percent, to ...</td>\n",
       "      <td>PG&amp;E Corp. shares jumped $1.63 or 8 percent to...</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Unnamed: 0                                                 s1  \\\n",
       "0           0  Amrozi accused his brother, whom he called \"th...   \n",
       "1           1  Yucaipa owned Dominick's before selling the ch...   \n",
       "2           2  They had published an advertisement on the Int...   \n",
       "3           3  Around 0335 GMT, Tab shares were up 19 cents, ...   \n",
       "4           4  The stock rose $2.11, or about 11 percent, to ...   \n",
       "\n",
       "                                                  s2  label corrected_label  \n",
       "0  Referring to him as only \"the witness\", Amrozi...      1            None  \n",
       "1  Yucaipa bought Dominick's in 1995 for $693 mil...      0            None  \n",
       "2  On June 10, the ship's owners had published an...      1            None  \n",
       "3  Tab shares jumped 20 cents, or 4.6%, to set a ...      0            None  \n",
       "4  PG&E Corp. shares jumped $1.63 or 8 percent to...      1            None  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split = \"train\"\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": "96c02b33",
   "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": "0628c592",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4076/4076 [05:22<00:00, 12.62it/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",
    "    total += 1\n",
    "    if row[\"label\"] == 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": "1ba89b2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4076\n",
      "1950\n",
      "0.1669802445907808\n",
      "355\n"
     ]
    }
   ],
   "source": [
    "print(total)\n",
    "print(no_need_correct)\n",
    "print(corrected/(corrected+cant_correct))\n",
    "print(corrected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "743357fb",
   "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": "8dab4e19",
   "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": "fa473aa0",
   "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>Amrozi accused his brother, whom he called \"th...</td>\n",
       "      <td>Referring to him as only \"the witness\", Amrozi...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Yucaipa owned Dominick's before selling the ch...</td>\n",
       "      <td>Yucaipa bought Dominick's in 1995 for $693 mil...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>They had published an advertisement on the Int...</td>\n",
       "      <td>On June 10, the ship's owners had published an...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Around 0335 GMT, Tab shares were up 19 cents, ...</td>\n",
       "      <td>Tab shares jumped 20 cents, or 4.6%, to set a ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>The stock rose $2.11, or about 11 percent, to ...</td>\n",
       "      <td>PG&amp;E Corp. shares jumped $1.63 or 8 percent to...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>Revenue in the first quarter of the year dropp...</td>\n",
       "      <td>With the scandal hanging over Stewart's compan...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>The Nasdaq had a weekly gain of 17.27, or 1.2 ...</td>\n",
       "      <td>The tech-laced Nasdaq Composite .IXIC rallied ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>The DVD-CCA then appealed to the state Supreme...</td>\n",
       "      <td>The DVD CCA appealed that decision to the U.S....</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>That compared with $35.18 million, or 24 cents...</td>\n",
       "      <td>Earnings were affected by a non-recurring $8 m...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>He said the foodservice pie business doesn't f...</td>\n",
       "      <td>\"The foodservice pie business does not fit our...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  s1  \\\n",
       "0  Amrozi accused his brother, whom he called \"th...   \n",
       "1  Yucaipa owned Dominick's before selling the ch...   \n",
       "2  They had published an advertisement on the Int...   \n",
       "3  Around 0335 GMT, Tab shares were up 19 cents, ...   \n",
       "4  The stock rose $2.11, or about 11 percent, to ...   \n",
       "5  Revenue in the first quarter of the year dropp...   \n",
       "6  The Nasdaq had a weekly gain of 17.27, or 1.2 ...   \n",
       "7  The DVD-CCA then appealed to the state Supreme...   \n",
       "8  That compared with $35.18 million, or 24 cents...   \n",
       "9  He said the foodservice pie business doesn't f...   \n",
       "\n",
       "                                                  s2  label  \n",
       "0  Referring to him as only \"the witness\", Amrozi...      1  \n",
       "1  Yucaipa bought Dominick's in 1995 for $693 mil...      0  \n",
       "2  On June 10, the ship's owners had published an...      1  \n",
       "3  Tab shares jumped 20 cents, or 4.6%, to set a ...      0  \n",
       "4  PG&E Corp. shares jumped $1.63 or 8 percent to...      0  \n",
       "5  With the scandal hanging over Stewart's compan...      0  \n",
       "6  The tech-laced Nasdaq Composite .IXIC rallied ...      0  \n",
       "7  The DVD CCA appealed that decision to the U.S....      0  \n",
       "8  Earnings were affected by a non-recurring $8 m...      0  \n",
       "9  \"The foodservice pie business does not fit our...      1  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_out.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3b8c4a05",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_out.to_csv(\"./mrpc_\"+split+\"_corrected.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65f5da03",
   "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
}
