second_looks.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab4efe82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.second_looks_utils import *\n",
    "from utils.train_utils import *\n",
    "import spacy\n",
    "import json\n",
    "seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a44d37",
   "metadata": {},
   "source": [
    "## Our reformatting functions\n",
    "- Redundant token removal\n",
    "- Example concatenation for longer training sequences\n",
    "- Preposing\n",
    "- Preposing + Interjection\n",
    "- ReCOGS\n",
    "- Variable-free"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68c03764",
   "metadata": {},
   "source": [
    "### Redundant Token Removal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80dda3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for removing_set in [\n",
    "    ['x', '_'],\n",
    "    ['x', '_', '(', ')'],\n",
    "    ['x', '_', '(', ',', ')']\n",
    "]:\n",
    "    def token_removal(text, phi): \n",
    "        # Parsing:\n",
    "        terms = []\n",
    "        for t in phi.split():\n",
    "            if t not in removing_set:\n",
    "                terms += [t]\n",
    "        ret = \" \".join(terms).strip()\n",
    "        return ret\n",
    "    train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "    train_df['LF'] = train_df[['sentence', 'LF']].apply(lambda x: token_removal(*x), axis=1,)\n",
    "    dev_df['LF'] = dev_df[['sentence', 'LF']].apply(lambda x: token_removal(*x), axis=1,)\n",
    "    test_df['LF'] = test_df[['sentence', 'LF']].apply(lambda x: token_removal(*x), axis=1,)\n",
    "    gen_df['LF'] = gen_df[['sentence', 'LF']].apply(lambda x: token_removal(*x), axis=1,)\n",
    "    \n",
    "    removing_set_str = \"\".join(removing_set)\n",
    "    dataset_postfix = f\"remove_{removing_set_str}\"\n",
    "    train_df.to_csv(f'./cogs_token_removal/train_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    dev_df.to_csv(f'./cogs_token_removal/dev_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    test_df.to_csv(f'./cogs_token_removal/test_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    gen_df.to_csv(f'./cogs_token_removal/gen_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6c1d099",
   "metadata": {},
   "source": [
    "### Example Concatenations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92cabc75",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reindex(LFs, initial_indexes):\n",
    "    new_LF_prefix = []\n",
    "    new_LF_body = []\n",
    "    for i in range(len(LFs)):\n",
    "        if initial_indexes[i] != 0:\n",
    "            new_lf = []\n",
    "            for item in LFs[i].split():\n",
    "                if item.isnumeric():\n",
    "                    new_i = int(item) + initial_indexes[i]\n",
    "                    new_lf += [str(new_i)]\n",
    "                else:\n",
    "                    new_lf += [item]\n",
    "            new_lf = \" \".join(new_lf)\n",
    "        else:\n",
    "            new_lf = LFs[i]\n",
    "        \n",
    "        for item in new_lf.split(\" ; \"):\n",
    "            if \"*\" in item:\n",
    "                new_LF_prefix += [item]\n",
    "            else:\n",
    "                new_LF_body += [item]\n",
    "        new_LF_body += [\"AND\"]\n",
    "    new_LF_body = new_LF_body[:-1]\n",
    "    return \" ; \".join(new_LF_prefix) + \" ; \" + \" \".join(new_LF_body)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7785a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "append_ks = [256, 512, 1024, 2048, 3072]\n",
    "for append_k in append_ks:\n",
    "    train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    train_df_org = train_df.copy()\n",
    "    train_df = train_df[train_df[\"type\"] != \"primitive\"]\n",
    "    dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "    dataset_postfix = f\"k_{append_k}\"\n",
    "    append_data = []\n",
    "    start_indexes = [i*6 for i in range(append_k)]\n",
    "    sorted_train_df = train_df.sort_values(by=\"sentence\", key=lambda x: x.str.len())\n",
    "    for start_index in start_indexes:\n",
    "        conj_1 = sorted_train_df.iloc[-2-start_index].sentence\n",
    "        if conj_1.split()[0] in {'The', 'A'}:\n",
    "            conj_1_first = conj_1[0].lower()\n",
    "        else:\n",
    "            conj_1_first = conj_1[0]\n",
    "            \n",
    "        conj_2 = sorted_train_df.iloc[-3-start_index].sentence\n",
    "        if conj_2.split()[0] in {'The', 'A'}:\n",
    "            conj_2_first = conj_2[0].lower()\n",
    "        else:\n",
    "            conj_2_first = conj_2[0]\n",
    "            \n",
    "        append_data += [\n",
    "            [sorted_train_df.iloc[-1-start_index].sentence[:-1]+\\\n",
    "            conj_1_first+\\\n",
    "            sorted_train_df.iloc[-2-start_index].sentence[1:-1]+\\\n",
    "            conj_2_first+\\\n",
    "            sorted_train_df.iloc[-3-start_index].sentence[1:],\n",
    "            reindex(\n",
    "                [\n",
    "                    sorted_train_df.iloc[-1-start_index].LF,\n",
    "                    sorted_train_df.iloc[-2-start_index].LF,\n",
    "                    sorted_train_df.iloc[-3-start_index].LF\n",
    "                ],\n",
    "                [\n",
    "                    0,\n",
    "                    len(sorted_train_df.iloc[-1-start_index].sentence[:-1].strip().split()),\n",
    "                    len(sorted_train_df.iloc[-1-start_index].sentence[:-1].strip().split())+\n",
    "                    len(sorted_train_df.iloc[-2-start_index].sentence[:-1].strip().split())\n",
    "                ]\n",
    "            ),\n",
    "            'concat']\n",
    "        ]\n",
    "    append_df = pd.DataFrame(append_data, columns =['sentence', 'LF', 'type'])\n",
    "    train_df = pd.concat([train_df_org, append_df])\n",
    "    train_df.to_csv(f'./cogs_concat/train_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    dev_df.to_csv(f'./cogs_concat/dev_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    test_df.to_csv(f'./cogs_concat/test_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    gen_df.to_csv(f'./cogs_concat/gen_{dataset_postfix}.tsv', sep='\\t', index=False, header=False)\n",
    "    \n",
    "    max_s = max(train_df['sentence'].str.split().apply(len))\n",
    "    max_lf = max(train_df['LF'].str.split().apply(len))\n",
    "    print(max_s, max_lf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ed5eb6f",
   "metadata": {},
   "source": [
    "### Preposing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845703e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "        \n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if \"nmod\" not in terms.split():\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= proposing_prob:\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if terms.split().count(\"nmod\") == 2:\n",
    "        upper_bound = 6\n",
    "        return_type = \"preposing_2\"\n",
    "    elif terms.split().count(\"nmod\") == 1:\n",
    "        upper_bound = 3\n",
    "        return_type = \"preposing_1\"\n",
    "    else:\n",
    "        assert False\n",
    "        \n",
    "    nmod = terms.split()[terms.split().index(\"nmod\")+2]\n",
    "    pre_phrase = text.split()[text.split().index(nmod)-2 : text.split().index(nmod)+upper_bound]\n",
    "    pre_phrase[0] = pre_phrase[0].capitalize()\n",
    "    pre_text = text.split()[:text.split().index(nmod)-2]\n",
    "    if pre_text[0] in [\"The\", \"A\"]:\n",
    "        pre_text[0] = pre_text[0].lower()\n",
    "    post_text = text.split()[text.split().index(nmod)+upper_bound:]\n",
    "    pre_text = pre_phrase + pre_text + post_text\n",
    "    pre_text = \" \".join(pre_text)\n",
    "\n",
    "    index_map = {}\n",
    "    idx = 0\n",
    "    for i in range(text.split().index(nmod)-2, text.split().index(nmod)+upper_bound):\n",
    "        index_map[f\"{i}\"] = f\"{idx}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)-2):\n",
    "        ii = 2+upper_bound+i\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)+upper_bound, len(text.split())):\n",
    "        ii = idx+(i-(text.split().index(nmod)+upper_bound))\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"  \n",
    "            \n",
    "    # now handle LF\n",
    "    pre_terms = []\n",
    "    for t in terms.split():\n",
    "        if t.isnumeric():\n",
    "            pre_terms += [str(index_map[str(int(t))])]\n",
    "        else:\n",
    "            pre_terms += [t]\n",
    "    pre_terms = \" \".join(pre_terms)\n",
    "\n",
    "    pre_terms_def = pre_terms.split(\" ; \")[:-1]\n",
    "    pre_terms_def.sort(key = lambda x: int(x.split()[-2]))  \n",
    "    pre_terms_role = pre_terms.split(\" ; \")[-1].split(\" AND \")\n",
    "    pre_terms_role.sort(key = lambda x: (int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else pre_text.split().index(x.split()[-2])))  \n",
    "    pre_terms_role = \" AND \".join(pre_terms_role)\n",
    "    pre_terms = \" ; \".join(pre_terms_def + [pre_terms_role])\n",
    "\n",
    "    return pre_text, pre_terms, return_type\n",
    "\n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "proposing_prob = 0.05\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00d26b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_postfix = \"preposing\"\n",
    "train_df.to_csv(f'./cogs_{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./cogs_{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./cogs_{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./cogs_{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96411af9",
   "metadata": {},
   "source": [
    "### Preposing + Sprinkles (Interjection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6f5c65d",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "        \n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if \"nmod\" not in terms.split():\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= proposing_prob:\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if terms.split().count(\"nmod\") == 2:\n",
    "        upper_bound = 6\n",
    "        return_type = \"preposing_2\"\n",
    "    elif terms.split().count(\"nmod\") == 1:\n",
    "        upper_bound = 3\n",
    "        return_type = \"preposing_1\"\n",
    "    else:\n",
    "        assert False\n",
    "        \n",
    "    nmod = terms.split()[terms.split().index(\"nmod\")+2]\n",
    "    pre_phrase = text.split()[text.split().index(nmod)-2 : text.split().index(nmod)+upper_bound]\n",
    "    pre_phrase[0] = pre_phrase[0].capitalize()\n",
    "    pre_text = text.split()[:text.split().index(nmod)-2]\n",
    "    if pre_text[0] in [\"The\", \"A\"]:\n",
    "        pre_text[0] = pre_text[0].lower()\n",
    "    post_text = text.split()[text.split().index(nmod)+upper_bound:]\n",
    "    pre_text = pre_phrase + pre_text + post_text\n",
    "    pre_text = \" \".join(pre_text)\n",
    "\n",
    "    index_map = {}\n",
    "    idx = 0\n",
    "    for i in range(text.split().index(nmod)-2, text.split().index(nmod)+upper_bound):\n",
    "        index_map[f\"{i}\"] = f\"{idx}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)-2):\n",
    "        ii = 2+upper_bound+i\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)+upper_bound, len(text.split())):\n",
    "        ii = idx+(i-(text.split().index(nmod)+upper_bound))\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"  \n",
    "            \n",
    "    # now handle LF\n",
    "    pre_terms = []\n",
    "    for t in terms.split():\n",
    "        if t.isnumeric():\n",
    "            pre_terms += [str(index_map[str(int(t))])]\n",
    "        else:\n",
    "            pre_terms += [t]\n",
    "    pre_terms = \" \".join(pre_terms)\n",
    "\n",
    "    pre_terms_def = pre_terms.split(\" ; \")[:-1]\n",
    "    pre_terms_def.sort(key = lambda x: int(x.split()[-2]))  \n",
    "    pre_terms_role = pre_terms.split(\" ; \")[-1].split(\" AND \")\n",
    "    pre_terms_role.sort(key = lambda x: (int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else pre_text.split().index(x.split()[-2])))  \n",
    "    pre_terms_role = \" AND \".join(pre_terms_role)\n",
    "    pre_terms = \" ; \".join(pre_terms_def + [pre_terms_role])\n",
    "\n",
    "    return pre_text, pre_terms, return_type\n",
    "\n",
    "def add_um(sentence):\n",
    "    words = sentence.split()\n",
    "    new_words = []\n",
    "    mapping = {}\n",
    "    offset = 0\n",
    "    for i, word in enumerate(words):\n",
    "        mapping[i] = len(new_words)\n",
    "        new_words.append(word)\n",
    "        if i > 0 and i < len(words) - 2 and random.random() > 0.5:\n",
    "            num_um = random.choice([1,2,3])\n",
    "            for j in range(num_um):\n",
    "                new_words.append(\"um\")\n",
    "    return \" \".join(new_words), mapping\n",
    "\n",
    "def sprinkle(text, phi, _type):\n",
    "    if \"preposition\" in _type:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    if random.random() >= sprinkle_prob:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    um_text, token_mapping = add_um(text)\n",
    "    um_phi = []\n",
    "    for t in phi.split():\n",
    "        if t.isnumeric():\n",
    "            um_phi += [str(token_mapping[int(t)])]\n",
    "        else:\n",
    "            um_phi += [t]\n",
    "    um_phi = \" \".join(um_phi)\n",
    "    \n",
    "    return um_text, um_phi, \"sprinkle\"\n",
    "            \n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "proposing_prob = 0.05\n",
    "sprinkle_prob = 0.05\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: sprinkle(*x), axis=1, result_type='expand')\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f2242b",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_postfix = \"preposing+sprinkles\"\n",
    "train_df.to_csv(f'./cogs_{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./cogs_{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./cogs_{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./cogs_{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d755a5d7",
   "metadata": {},
   "source": [
    "### ReCOGS (Number of resampling iterations = 5)\n",
    "\n",
    "It seems like the performance gain from increasing the number of resampling iterations dimish quickly after getting the number above 10. We are trying 5 here.\n",
    "\n",
    "**Since the first release of ReCOGS, we made some changes to the dataset:**\n",
    "\n",
    "- NP and PP phrases mirror the word orders, instead of regrouping together in the output logical forms.\n",
    "- Adding in examples with preposing and interjection in the training dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9435232",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "        \n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if \"nmod\" not in terms.split():\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= proposing_prob:\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if terms.split().count(\"nmod\") == 2:\n",
    "        upper_bound = 6\n",
    "        return_type = \"preposing_2\"\n",
    "    elif terms.split().count(\"nmod\") == 1:\n",
    "        upper_bound = 3\n",
    "        return_type = \"preposing_1\"\n",
    "    else:\n",
    "        assert False\n",
    "        \n",
    "    nmod = terms.split()[terms.split().index(\"nmod\")+2]\n",
    "    pre_phrase = text.split()[text.split().index(nmod)-2 : text.split().index(nmod)+upper_bound]\n",
    "    pre_phrase[0] = pre_phrase[0].capitalize()\n",
    "    pre_text = text.split()[:text.split().index(nmod)-2]\n",
    "    if pre_text[0] in [\"The\", \"A\"]:\n",
    "        pre_text[0] = pre_text[0].lower()\n",
    "    post_text = text.split()[text.split().index(nmod)+upper_bound:]\n",
    "    pre_text = pre_phrase + pre_text + post_text\n",
    "    pre_text = \" \".join(pre_text)\n",
    "\n",
    "    index_map = {}\n",
    "    idx = 0\n",
    "    for i in range(text.split().index(nmod)-2, text.split().index(nmod)+upper_bound):\n",
    "        index_map[f\"{i}\"] = f\"{idx}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)-2):\n",
    "        ii = 2+upper_bound+i\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)+upper_bound, len(text.split())):\n",
    "        ii = idx+(i-(text.split().index(nmod)+upper_bound))\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"  \n",
    "            \n",
    "    # now handle LF\n",
    "    pre_terms = []\n",
    "    for t in terms.split():\n",
    "        if t.isnumeric():\n",
    "            pre_terms += [str(index_map[str(int(t))])]\n",
    "        else:\n",
    "            pre_terms += [t]\n",
    "    pre_terms = \" \".join(pre_terms)\n",
    "\n",
    "    pre_terms_def = pre_terms.split(\" ; \")[:-1]\n",
    "    pre_terms_def.sort(key = lambda x: int(x.split()[-2]))  \n",
    "    pre_terms_role = pre_terms.split(\" ; \")[-1].split(\" AND \")\n",
    "    pre_terms_role.sort(key = lambda x: (int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else pre_text.split().index(x.split()[-2])))  \n",
    "    pre_terms_role = \" AND \".join(pre_terms_role)\n",
    "    pre_terms = \" ; \".join(pre_terms_def + [pre_terms_role])\n",
    "\n",
    "    return pre_text, pre_terms, return_type\n",
    "\n",
    "def add_um(sentence):\n",
    "    words = sentence.split()\n",
    "    new_words = []\n",
    "    mapping = {}\n",
    "    offset = 0\n",
    "    for i, word in enumerate(words):\n",
    "        mapping[i] = len(new_words)\n",
    "        new_words.append(word)\n",
    "        if i > 0 and i < len(words) - 2 and random.random() > 0.5:\n",
    "            num_um = random.choice([1,2,3])\n",
    "            for j in range(num_um):\n",
    "                new_words.append(\"um\")\n",
    "    return \" \".join(new_words), mapping\n",
    "\n",
    "def sprinkle(text, phi, _type):\n",
    "    if \"preposition\" in _type:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    if random.random() >= sprinkle_prob:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    um_text, token_mapping = add_um(text)\n",
    "    um_phi = []\n",
    "    for t in phi.split():\n",
    "        if t.isnumeric():\n",
    "            um_phi += [str(token_mapping[int(t)])]\n",
    "        else:\n",
    "            um_phi += [t]\n",
    "    um_phi = \" \".join(um_phi)\n",
    "    \n",
    "    return um_text, um_phi, f\"{_type}+sprinkle\"\n",
    "            \n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "proposing_prob = 0.05\n",
    "sprinkle_prob = 0.05\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: sprinkle(*x), axis=1, result_type='expand')\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2541f2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_digit_pool = set([])\n",
    "# loading target vocab to random sample our variable names\n",
    "for k, v in load_vocab(\"./model/tgt_vocab.txt\").items():\n",
    "    if k.isnumeric():\n",
    "        existing_digit_pool.add(k)\n",
    "existing_digit_pool = list(existing_digit_pool)\n",
    "\n",
    "def translate(text, phi):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, f\"LAMBDA a . {phi} ( a )\"\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        if len(phi.split()) == 7:\n",
    "            return text, phi\n",
    "        phi_split = phi.split(text)\n",
    "        cleaned_phi = []\n",
    "        for chunk in phi_split:\n",
    "            if \"LAMBDA\" in chunk:\n",
    "                cleaned_phi += [chunk.strip()]\n",
    "            else:\n",
    "                verb_args = chunk.strip(\" .\").split()[2]\n",
    "                cleaned_phi += [chunk.strip(\" .\")]\n",
    "        return text, \" \".join(cleaned_phi[:1] + [f\"{text} ( {verb_args} ) AND\"] + cleaned_phi[1:])\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "            if \"x _\" not in d['entvar']:\n",
    "                d['entvar_name'] = d['entvar']\n",
    "                assert text_split.count(d['entvar']) == 1\n",
    "                name_idx = text_split.index(d['entvar'])\n",
    "                d['entvar'] = f\"x _ {name_idx}\"\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                def_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            if f\"{d['pred']} ( {d['eventvar']} )\" not in role_terms:\n",
    "                role_terms += [f\"{d['pred']} ( {d['eventvar']} )\"]\n",
    "            role_terms += [f\"{d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "            if \"entvar_name\" in d:\n",
    "                def_terms += [f\"{d['entvar_name']} ( {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            role_terms += [f\"nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))\n",
    "\n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    # final step, remove biases\n",
    "    current_digit_pool = set([])\n",
    "    for t in terms.split():\n",
    "        if t.isnumeric():\n",
    "            current_digit_pool.add(t)\n",
    "    current_digit_pool = list(current_digit_pool)\n",
    "    random.shuffle(current_digit_pool)\n",
    "    sample_random_digit = random.sample(existing_digit_pool, k=len(current_digit_pool))\n",
    "    digit_mapping = dict(zip(current_digit_pool, sample_random_digit))\n",
    "\n",
    "    new_terms = []\n",
    "    for t in terms.split():\n",
    "        if t == \"_\" or t == \"x\":\n",
    "            continue\n",
    "        if t.isnumeric():\n",
    "            new_terms += [digit_mapping[t]]\n",
    "        else:\n",
    "            new_terms += [t]\n",
    "\n",
    "    terms = \" \".join(new_terms)\n",
    "    return text, terms\n",
    "\n",
    "sampled_n = 5\n",
    "append_k = 3072\n",
    "\n",
    "train_dfs = []\n",
    "for i in range(sampled_n):\n",
    "    train_df_i = train_df.copy()\n",
    "    train_df_i[['sentence', 'LF']] = train_df_i[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "    train_dfs += [train_df_i]\n",
    "dev_df[['sentence', 'LF']] = dev_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF']] = test_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF']] = gen_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "\n",
    "def reindex(LFs, existing_digit_pool):\n",
    "    curr_digit = set([])\n",
    "    for i in range(len(LFs)):\n",
    "        for item in LFs[i].split():\n",
    "            if item.isnumeric():\n",
    "                curr_digit.add((i, int(item)))\n",
    "    sampled_digits = random.sample(existing_digit_pool, k=len(curr_digit))\n",
    "    digit_map = {}\n",
    "    idx = 0\n",
    "    for d in list(curr_digit):\n",
    "        digit_map[d] = sampled_digits[idx]\n",
    "        idx += 1\n",
    "    \n",
    "    reindex_LFs = []\n",
    "    for i in range(len(LFs)):\n",
    "        new_LFs = []\n",
    "        for item in LFs[i].split():\n",
    "            if item.isnumeric():\n",
    "                new_LFs += [digit_map[(i, int(item))]]\n",
    "            else:\n",
    "                new_LFs += [item]\n",
    "        reindex_LFs += [\" \".join(new_LFs)]\n",
    "        \n",
    "    new_LF_prefix = []\n",
    "    new_LF_body_role = []\n",
    "        \n",
    "    for i in range(len(reindex_LFs)):\n",
    "        new_LF_prefix.extend(reindex_LFs[i].split(\" ; \")[:-1])\n",
    "        for term in reindex_LFs[i].split(\" ; \")[-1].split(\" AND \"):\n",
    "            new_LF_body_role += [term]\n",
    "                \n",
    "    new_LF_body = new_LF_body_role\n",
    "        \n",
    "    return \" ; \".join(new_LF_prefix) + \" ; \" + \" AND \".join(new_LF_body)\n",
    "\n",
    "start_indexes = [i*6 for i in range(append_k)]\n",
    "append_data = []\n",
    "\n",
    "for i in range(sampled_n):\n",
    "    train_df_sorted = train_dfs[i].sort_values(by=\"sentence\", key=lambda x: x.str.len())\n",
    "    for start_index in start_indexes:\n",
    "        conj_1 = train_df_sorted.iloc[-2-start_index].sentence\n",
    "        if conj_1.split()[0] in {'The', 'A'}:\n",
    "            conj_1_first = conj_1[0].lower()\n",
    "        else:\n",
    "            conj_1_first = conj_1[0]\n",
    "            \n",
    "        conj_2 = train_df_sorted.iloc[-3-start_index].sentence\n",
    "        if conj_2.split()[0] in {'The', 'A'}:\n",
    "            conj_2_first = conj_2[0].lower()\n",
    "        else:\n",
    "            conj_2_first = conj_2[0]\n",
    "            \n",
    "        append_data += [\n",
    "            [train_df_sorted.iloc[-1-start_index].sentence[:-1]+\\\n",
    "            conj_1_first+\\\n",
    "            train_df_sorted.iloc[-2-start_index].sentence[1:-1]+\\\n",
    "            conj_2_first+\\\n",
    "            train_df_sorted.iloc[-3-start_index].sentence[1:],\n",
    "            reindex(\n",
    "                [\n",
    "                    train_df_sorted.iloc[-1-start_index].LF,\n",
    "                    train_df_sorted.iloc[-2-start_index].LF,\n",
    "                    train_df_sorted.iloc[-3-start_index].LF\n",
    "                ], existing_digit_pool\n",
    "            ),\n",
    "            'length_ood']\n",
    "        ]\n",
    "append_data = pd.DataFrame(append_data, columns =['sentence', 'LF', 'type'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471bc3a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.concat(train_dfs)\n",
    "train_df = pd.concat([train_df, append_data])\n",
    "train_df = train_df.drop_duplicates()\n",
    "\n",
    "dataset_postfix = \"recogs_v2\"\n",
    "train_df.to_csv(f'./{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "657a5706",
   "metadata": {},
   "source": [
    "ReCOGS testing function to translate back to the COGS LF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b53f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "verb_to_forms_map = {}\n",
    "def collect_verb_forms(text, phi):\n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "            verb = d['pred']\n",
    "            if \"x _\" in d['eventvar']:\n",
    "                form = text.split()[int(d['eventvar'].split()[-1])]\n",
    "                if verb not in verb_to_forms_map:\n",
    "                    verb_to_forms_map[verb] = set([form])\n",
    "                else:\n",
    "                    verb_to_forms_map[verb].add(form)\n",
    "\n",
    "_ = train_df[['sentence', 'LF']].apply(lambda x: collect_verb_forms(*x), axis=1, result_type='expand')\n",
    "_ = dev_df[['sentence', 'LF']].apply(lambda x: collect_verb_forms(*x), axis=1, result_type='expand')\n",
    "_ = test_df[['sentence', 'LF']].apply(lambda x: collect_verb_forms(*x), axis=1, result_type='expand')\n",
    "_ = gen_df[['sentence', 'LF']].apply(lambda x: collect_verb_forms(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aff6f31",
   "metadata": {},
   "outputs": [],
   "source": [
    "recogs_neoD_verb_re = re.compile(r\"\"\"\n",
    "    ^\n",
    "    \\s*(\\w+?)\\s*\n",
    "    \\(\n",
    "    \\s*([0-9]+?)\\s*\n",
    "    \\)\n",
    "    \\s*$\"\"\", re.VERBOSE)\n",
    "\n",
    "recogs_neoD_pred_re = re.compile(r\"\"\"\n",
    "    ^\n",
    "    \\s*(\\w+?)\\s*\n",
    "    \\(\n",
    "    \\s*(.+?)\\s*\n",
    "    ,\n",
    "    \\s*(.+?)\\s*\n",
    "    \\)\n",
    "    \\s*$\"\"\", re.VERBOSE)\n",
    "\n",
    "recogs_neoD_mod_re = re.compile(r\"\"\"\n",
    "    ^\n",
    "    \\s*(\\w+?)\\s*\n",
    "    \\.\n",
    "    \\s*(\\w+?)\\s*\n",
    "    \\(\n",
    "    \\s*(.+?)\\s*\n",
    "    ,\n",
    "    \\s*(.+?)\\s*\n",
    "    \\)\n",
    "    \\s*$\"\"\", re.VERBOSE)\n",
    "\n",
    "def reverse_translate(text, phi):\n",
    "    conjs = phi.split(\" ; \")[-1]\n",
    "    nouns = phi.split(\" ; \")[:-1]\n",
    "    nouns_map = {}\n",
    "    text_split = text.split()\n",
    "    filtered_nouns = []\n",
    "    noun_terms = []\n",
    "    for noun in nouns:\n",
    "        if noun.split()[0][0].isupper():\n",
    "            nouns_map[noun.split()[-2]] = (noun.split()[0], )\n",
    "        else:\n",
    "            filtered_nouns += [noun]\n",
    "            if \"*\" in noun:\n",
    "                noun_str = noun.split()[1]                \n",
    "            else:\n",
    "                noun_str = noun.split()[0]\n",
    "            if noun_str in {'tv'}:\n",
    "                nouns_map[noun.split()[-2]] = (str(text_split.index('TV')), noun_str)\n",
    "            else:\n",
    "                nouns_map[noun.split()[-2]] = (str(text_split.index(noun_str)), noun_str)\n",
    "    \n",
    "    filtered_conjs = []\n",
    "    for noun in filtered_nouns:\n",
    "        if np_re.search(noun):\n",
    "            d = parse_np(noun)\n",
    "            if \"*\" in noun:\n",
    "                noun_terms += [f\"* {d['pred']} ( x _ {nouns_map[d['entvar']][0]} )\"]\n",
    "            else:\n",
    "                filtered_conjs += [f\"{d['pred']} ( x _ {nouns_map[d['entvar']][0]} )\"]\n",
    "      \n",
    "    verb_map = {}\n",
    "    for conj in conjs.split(\" AND \"):\n",
    "        if recogs_neoD_verb_re.search(conj):\n",
    "            # candidate for mapping verb.\n",
    "            pred, arg = recogs_neoD_verb_re.search(conj).groups()\n",
    "            verb_map[arg] = pred\n",
    "                \n",
    "    for conj in conjs.split(\" AND \"):\n",
    "        if \"nmod\" in conj:\n",
    "            role, pred, first_arg, second_arg = recogs_neoD_mod_re.search(conj).groups()\n",
    "            filtered_conjs += [f\"{nouns_map[first_arg][1]} . nmod . {pred} ( x _ {nouns_map[first_arg][0]} , x _ {nouns_map[second_arg][0]} )\"]\n",
    "        else:\n",
    "            if recogs_neoD_verb_re.search(conj):\n",
    "                pass\n",
    "            else:\n",
    "                role, first_arg, second_arg = recogs_neoD_pred_re.search(conj).groups()\n",
    "                forms_set = verb_to_forms_map[verb_map[first_arg]]\n",
    "                verb_idx = -1\n",
    "                for form in list(forms_set):\n",
    "                    if form in text_split:\n",
    "                        verb_idx = str(text_split.index(form))\n",
    "                        break\n",
    "                if second_arg in nouns_map:\n",
    "                    if len(nouns_map[second_arg]) == 1:\n",
    "                        filtered_conjs += [f\"{verb_map[first_arg]} . {role} ( x _ {verb_idx} , {nouns_map[second_arg][0]} )\"]\n",
    "                    else:\n",
    "                        filtered_conjs += [f\"{verb_map[first_arg]} . {role} ( x _ {verb_idx} , x _ {nouns_map[second_arg][0]} )\"]\n",
    "                else:\n",
    "                    forms_set = verb_to_forms_map[verb_map[second_arg]]\n",
    "                    comp_verb_idx = -1\n",
    "                    for form in list(forms_set):\n",
    "                        if form in text_split:\n",
    "                            comp_verb_idx = str(text_split.index(form))\n",
    "                            break\n",
    "                    filtered_conjs += [f\"{verb_map[first_arg]} . {role} ( x _ {verb_idx} , x _ {comp_verb_idx} )\"]\n",
    "    filtered_conjs.sort(key = lambda x: (int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else text_split.index(x.split()[-2])))  \n",
    "    \n",
    "    if len(noun_terms) > 0:\n",
    "        lf = \" ; \".join(noun_terms) + \" ; \" + \" AND \".join(filtered_conjs)\n",
    "    else:\n",
    "        lf = \" AND \".join(filtered_conjs)\n",
    "    return phi, lf\n",
    "    \n",
    "recogs_dev_df = pd.read_csv(\"./recogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "recogs_test_df = pd.read_csv(\"./recogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "recogs_gen_df = pd.read_csv(\"./recogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "recogs_dev_df[['recogs_LF', 'backward_LF']] = recogs_dev_df[['sentence', 'LF']].apply(\n",
    "    lambda x: reverse_translate(*x), axis=1, result_type='expand'\n",
    ")\n",
    "recogs_test_df[['recogs_LF', 'backward_LF']] = recogs_test_df[['sentence', 'LF']].apply(\n",
    "    lambda x: reverse_translate(*x), axis=1, result_type='expand'\n",
    ")\n",
    "recogs_gen_df[['recogs_LF', 'backward_LF']] = recogs_gen_df[['sentence', 'LF']].apply(\n",
    "    lambda x: reverse_translate(*x), axis=1, result_type='expand'\n",
    ")\n",
    "combined_dev_df = pd.concat([dev_df, recogs_dev_df[['recogs_LF', 'backward_LF']]], axis=1)\n",
    "combined_test_df = pd.concat([test_df, recogs_test_df[['recogs_LF', 'backward_LF']]], axis=1)\n",
    "combined_gen_df = pd.concat([gen_df, recogs_gen_df[['recogs_LF', 'backward_LF']]], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d08f9e9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: i am a bit lazy here... you can work out a script without needing this.\n",
    "df_inspecting = combined_gen_df\n",
    "print(\"WARNING: Be aware these errors are expected as ReCOGS focuses on semantics reprs not recoversing token positions!\")\n",
    "print()\n",
    "print(\"NUM OF ERRORS: \", len(df_inspecting[df_inspecting['LF'] != df_inspecting['backward_LF']]))\n",
    "print()\n",
    "inspect_id = 0\n",
    "row = df_inspecting[df_inspecting['LF'] != df_inspecting['backward_LF']].iloc[inspect_id]\n",
    "print(\"sentence: \", row['sentence'])\n",
    "print()\n",
    "print(\"LF: \", row['LF'])\n",
    "print()\n",
    "print(\"backward_LF: \", row['backward_LF'])\n",
    "print()\n",
    "print(\"recogs_LF: \", row['recogs_LF'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3bc584",
   "metadata": {},
   "source": [
    "### Variable-free format\n",
    "This is from Qiu et. al., 2022 paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e28d543b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cogs_lf_to_funcall(lf):\n",
    "    \"\"\"Converts the given COGS logical form into the variable-free form.\n",
    "    - Nouns (entities and unaries) become values:\n",
    "      Jack --> Jack\n",
    "      cat ( x _ 1 ) --> cat\n",
    "      * cat ( x _ 1 ) --> * cat\n",
    "    - Verbs become functions, and their roles become argument names:\n",
    "      eat . agent ( x _ 2 , Jack ) --> eat ( agent = Jack )\n",
    "    - The variables representing nouns resolve to their values:\n",
    "      cat ( x _ 1 ) AND eat . agent ( x _ 2 , x _ 1 ) --> eat ( agent = cat )\n",
    "    This converter constructs a graph where variables are nodes and binaries\n",
    "    are edges. After identifying the root, it then performs depth-first traversal\n",
    "    to construct the output.\n",
    "    Args:\n",
    "    lf: Logical form string.\n",
    "    Returns:\n",
    "    The converted logical form.\n",
    "    \"\"\"\n",
    "    if \"LAMBDA\" in lf or \"(\" not in lf:\n",
    "        return lf\n",
    "\n",
    "    # Parse the terms in the logical form\n",
    "    # Example: toss . agent ( x _ 1 , John ) --> [toss, agent], [x _ 1, John]\n",
    "    terms = []\n",
    "    for raw_term in re.split(\" ; | AND \", lf):\n",
    "        match = re.match(r\"(.*) \\( (.*) \\)\", raw_term)\n",
    "        if not match:\n",
    "            raise ValueError(f\"Malformed term: {raw_term}\")\n",
    "        labels = match.group(1).split(\" . \")\n",
    "        args = match.group(2).split(\" , \")\n",
    "        if len(args) not in (1, 2):\n",
    "            raise ValueError(f\"Invalid number of args: {args}\")\n",
    "        terms.append((labels, args))\n",
    "\n",
    "    # `nodes` maps variables to node name (e.g., \"x _ 3\" -> \"* cat\").\n",
    "    nodes = {}\n",
    "    for labels, args in terms:\n",
    "        if args[0] in nodes:\n",
    "            # The variable has already been seen; check for conflicts.\n",
    "            if nodes[args[0]] not in (labels[0], \"* \" + labels[0]):\n",
    "                raise ValueError(\n",
    "                    f\"Conflicting node name: {nodes[args[0]]} vs. {labels[0]}\")\n",
    "        else:\n",
    "            nodes[args[0]] = labels[0]\n",
    "\n",
    "    # `children` maps variables to a list of (edge name, target node).\n",
    "    children = {}\n",
    "    # Potential root nodes; any node being a child will be removed.\n",
    "    root_candidates = list(nodes)\n",
    "    for labels, args in terms:\n",
    "        if len(args) == 2:\n",
    "            if args[0] not in children:\n",
    "                children[args[0]] = []\n",
    "            children[args[0]].append((\" . \".join(labels[1:]), args[1]))\n",
    "            if args[1] in root_candidates:\n",
    "                root_candidates.remove(args[1])\n",
    "    if len(root_candidates) != 1:\n",
    "        raise ValueError(f\"Multiple roots: {root_candidates}\")\n",
    "    root = root_candidates[0]\n",
    "\n",
    "    # Depth-first traverse the graph to construct the funcall\n",
    "    def dfs(node):\n",
    "        if node not in nodes:\n",
    "            # Named entity such as \"John\"\n",
    "            if node.startswith(\"x _\"):\n",
    "                raise ValueError(f\"Unbound variable {node}\")\n",
    "            if node in children:\n",
    "                raise ValueError(f\"Named entity {node} has children {children[node]}\")\n",
    "            return [node]\n",
    "        else:\n",
    "            # A noun like \"cat\" or a verb like \"jump\"\n",
    "            if node not in children:\n",
    "                return [nodes[node]]\n",
    "            funcall_args = []\n",
    "            for edge_label, edge_target in children[node]:\n",
    "                funcall_args.append([edge_label, \"=\"] + dfs(edge_target))\n",
    "            funcall = [nodes[node], \"(\"]\n",
    "            for i, funcall_arg in enumerate(funcall_args):\n",
    "                if i != 0:\n",
    "                    funcall.append(\",\")\n",
    "                funcall.extend(funcall_arg)\n",
    "            funcall.append(\")\")\n",
    "            return funcall\n",
    "\n",
    "    return \" \".join(dfs(root))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec6c7fbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "train_df['LF'] = train_df[['LF']].apply(lambda x: cogs_lf_to_funcall(*x), axis=1,)\n",
    "dev_df['LF'] = dev_df[['LF']].apply(lambda x: cogs_lf_to_funcall(*x), axis=1,)\n",
    "test_df['LF'] = test_df[['LF']].apply(lambda x: cogs_lf_to_funcall(*x), axis=1,)\n",
    "gen_df['LF'] = gen_df[['LF']].apply(lambda x: cogs_lf_to_funcall(*x), axis=1,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72eb5e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_postfix = \"variable_free\"\n",
    "train_df.to_csv(f'./{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe3cd798",
   "metadata": {},
   "source": [
    "### Participial Verb Phrases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4548149",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "\n",
    "noun_pool = set()\n",
    "v_pool = set()\n",
    "def collect_nv(text, phi, _type):\n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    role_set = set([])\n",
    "    for d in data:\n",
    "        if d['type'] == 'role':\n",
    "            v_pool.add(d['pred'])\n",
    "        elif d['type'] == 'np':\n",
    "            noun_pool.add(d['pred'])\n",
    "            \n",
    "    return text, phi, _type\n",
    "\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "\n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    role_set = set([])\n",
    "    object_entvar = None\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "            if d['role'] == \"theme\":\n",
    "                object_entvar = d['entvar']\n",
    "                \n",
    "            role_set.add(d['role'])\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    start_with_pp = False if text_split[0] in {'A', 'The'} else True\n",
    "    \n",
    "    if start_with_pp or (object_entvar is not None and \"x _ \" not in object_entvar):\n",
    "        # combine\n",
    "        def_terms = \" ; \".join(def_terms)\n",
    "        if def_terms == \"\":\n",
    "            terms = \" AND \".join(rest_terms)\n",
    "        elif \" AND \".join(rest_terms) == \"\":\n",
    "            terms = def_terms\n",
    "        else:\n",
    "            terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= participial_prob:\n",
    "        # combine\n",
    "        def_terms = \" ; \".join(def_terms)\n",
    "        if def_terms == \"\":\n",
    "            terms = \" AND \".join(rest_terms)\n",
    "        elif \" AND \".join(rest_terms) == \"\":\n",
    "            terms = def_terms\n",
    "        else:\n",
    "            terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "        return text, terms, _type\n",
    "    \n",
    "    # agent -> subject\n",
    "    # theme -> object\n",
    "    if \"theme\" in phi:\n",
    "        # only one object if it exists!\n",
    "        assert 1 == phi.split().count(\"theme\")\n",
    "        modify_subject = True if random.random() < 0.5 else False\n",
    "    else:\n",
    "        modify_subject = True\n",
    "    \n",
    "    # reindex!\n",
    "    def reindex_terms(terms, offset=1, distance=3):\n",
    "        new_terms = []\n",
    "        for term in terms:\n",
    "            new_term = []\n",
    "            for t in term.split():\n",
    "                if t.isnumeric() and int(t) > offset:\n",
    "                    new_term += [str(int(t)+distance)]\n",
    "                else:\n",
    "                    new_term += [t]\n",
    "            new_terms += [\" \".join(new_term)]\n",
    "        return new_terms\n",
    "    \n",
    "    # we need to decide whether to add to subj or obj.\n",
    "    second_text_pp = False\n",
    "    if modify_subject:\n",
    "        subject = text_split[0] if start_with_pp else text_split[1]\n",
    "        verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "        noun_pp = random.choice(list(noun_pool))\n",
    "        articles = 'a' if random.random() < 0.5 else 'the'\n",
    "        text_pp = [verb_pp_map[verb_pp], articles, noun_pp]\n",
    "        if articles == 'a':\n",
    "            pp_np_term =  [f\"{noun_pp} ( x _ 4 )\"]\n",
    "        elif articles == 'the':\n",
    "            pp_np_term = [f\"* {noun_pp} ( x _ 4 )\"]\n",
    "        pp_role_term = [f\"{subject} . acl . {verb_pp} ( x _ 1 , x _ 4 )\"]\n",
    "        \n",
    "        if random.random() < 0.5:\n",
    "            second_text_pp = True\n",
    "            second_verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "            second_noun_pp = random.choice(list(noun_pool))\n",
    "            second_articles = 'a' if random.random() < 0.5 else 'the'\n",
    "            text_pp += [verb_pp_map[second_verb_pp], second_articles, second_noun_pp]\n",
    "            if second_articles == 'a':\n",
    "                pp_np_term += [f\"{second_noun_pp} ( x _ 7 )\"]\n",
    "            elif second_articles == 'the':\n",
    "                pp_np_term += [f\"* {second_noun_pp} ( x _ 7 )\"]\n",
    "            pp_role_term += [f\"{noun_pp} . acl . {second_verb_pp} ( x _ 4 , x _ 7 )\"]\n",
    "            \n",
    "        def_terms = reindex_terms(def_terms, distance=6 if second_text_pp else 3)\n",
    "        rest_terms = reindex_terms(rest_terms, distance=6 if second_text_pp else 3)\n",
    "        \n",
    "        object_pos = 1\n",
    "        \n",
    "    else:\n",
    "        # locate the obj\n",
    "        assert \"x _ \" in object_entvar\n",
    "        object_pos = int(object_entvar.split()[-1])\n",
    "        _object = text_split[object_pos]\n",
    "        \n",
    "        verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "        noun_pp = random.choice(list(noun_pool))\n",
    "        articles = 'a' if random.random() < 0.5 else 'the'\n",
    "        text_pp = [verb_pp_map[verb_pp], articles, noun_pp]\n",
    "        if articles == 'a':\n",
    "            pp_np_term = [f\"{noun_pp} ( x _ {object_pos+3} )\"]\n",
    "        elif articles == 'the':\n",
    "            pp_np_term = [f\"* {noun_pp} ( x _ {object_pos+3} )\"]\n",
    "        pp_role_term = [f\"{_object} . acl . {verb_pp} ( x _ {object_pos} , x _ {object_pos+3} )\"]\n",
    "        \n",
    "        if random.random() < 0.5:\n",
    "            second_text_pp = True\n",
    "            second_verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "            second_noun_pp = random.choice(list(noun_pool))\n",
    "            second_articles = 'a' if random.random() < 0.5 else 'the'\n",
    "            text_pp += [verb_pp_map[second_verb_pp], second_articles, second_noun_pp]\n",
    "            if second_articles == 'a':\n",
    "                pp_np_term += [f\"{second_noun_pp} ( x _ {object_pos+6} )\"]\n",
    "            elif second_articles == 'the':\n",
    "                pp_np_term += [f\"* {second_noun_pp} ( x _ {object_pos+6} )\"]\n",
    "            pp_role_term += [f\"{noun_pp} . acl . {second_verb_pp} ( x _ {object_pos+3} , x _ {object_pos+6} )\"]\n",
    " \n",
    "        def_terms = reindex_terms(def_terms, object_pos, distance=6 if second_text_pp else 3)\n",
    "        rest_terms = reindex_terms(rest_terms, object_pos, distance=6 if second_text_pp else 3)\n",
    "    \n",
    "    for np_t in pp_np_term:\n",
    "        if np_t[0] == \"*\":\n",
    "            def_terms += [np_t]\n",
    "        else:\n",
    "            rest_terms += [np_t]\n",
    "    rest_terms += pp_role_term\n",
    "    text_split = text_split[:object_pos+1] + text_pp + text_split[object_pos+1:]\n",
    "    new_text = \" \".join(text_split)\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))   \n",
    "    rest_terms.sort(\n",
    "        key = lambda x: (\n",
    "            int(x.split()[-7]) if x.split()[-1] == \"acl\" else int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), \n",
    "            int(x.split()[-3]) if x.split()[-1] == \"acl\" else -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else new_text.split().index(x.split()[-2])\n",
    "        )\n",
    "    )  \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if second_text_pp:\n",
    "        return new_text, terms, \"subj_participle_verb_phrase_nested\" if modify_subject else \"obj_participle_verb_phrase_nested\"\n",
    "    else:\n",
    "        return new_text, terms, \"subj_participle_verb_phrase\" if modify_subject else \"obj_participle_verb_phrase\"\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c9f375",
   "metadata": {},
   "outputs": [],
   "source": [
    "participial_prob = 0.15\n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "with open('./verb_pp_map.json', 'r') as openfile:\n",
    "    verb_pp_map = json.load(openfile)\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(\n",
    "    lambda x: collect_nv(*x), axis=1, result_type='expand'\n",
    ")\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(\n",
    "    lambda x: translate(*x), axis=1, result_type='expand'\n",
    ")\n",
    "\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d521dcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_postfix = \"participle_verb\"\n",
    "train_df.to_csv(f'./cogs_{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./cogs_{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./cogs_{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./cogs_{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f541db60",
   "metadata": {},
   "source": [
    "### Participial Verb Phrases (Easy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fd9e356",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "\n",
    "noun_pool = set()\n",
    "v_pool = set()\n",
    "def collect_nv(text, phi, _type):\n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    role_set = set([])\n",
    "    for d in data:\n",
    "        if d['type'] == 'role':\n",
    "            v_pool.add(d['pred'])\n",
    "        elif d['type'] == 'np':\n",
    "            noun_pool.add(d['pred'])\n",
    "            \n",
    "    return text, phi, _type\n",
    "\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "\n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    role_set = set([])\n",
    "    object_entvar = None\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "            if d['role'] == \"theme\":\n",
    "                object_entvar = d['entvar']\n",
    "                \n",
    "            role_set.add(d['role'])\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    start_with_pp = False if text_split[0] in {'A', 'The'} else True\n",
    "    \n",
    "    if start_with_pp or (object_entvar is not None and \"x _ \" not in object_entvar):\n",
    "        # combine\n",
    "        def_terms = \" ; \".join(def_terms)\n",
    "        if def_terms == \"\":\n",
    "            terms = \" AND \".join(rest_terms)\n",
    "        elif \" AND \".join(rest_terms) == \"\":\n",
    "            terms = def_terms\n",
    "        else:\n",
    "            terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= participial_prob:\n",
    "        # combine\n",
    "        def_terms = \" ; \".join(def_terms)\n",
    "        if def_terms == \"\":\n",
    "            terms = \" AND \".join(rest_terms)\n",
    "        elif \" AND \".join(rest_terms) == \"\":\n",
    "            terms = def_terms\n",
    "        else:\n",
    "            terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "        return text, terms, _type\n",
    "    \n",
    "    # agent -> subject\n",
    "    # theme -> object\n",
    "    if \"theme\" in phi:\n",
    "        # only one object if it exists!\n",
    "        assert 1 == phi.split().count(\"theme\")\n",
    "        modify_subject = True if random.random() < 0.5 else False\n",
    "    else:\n",
    "        modify_subject = True\n",
    "    \n",
    "    # reindex!\n",
    "    def reindex_terms(terms, offset=1, distance=3):\n",
    "        new_terms = []\n",
    "        for term in terms:\n",
    "            new_term = []\n",
    "            for t in term.split():\n",
    "                if t.isnumeric() and int(t) > offset:\n",
    "                    new_term += [str(int(t)+distance)]\n",
    "                else:\n",
    "                    new_term += [t]\n",
    "            new_terms += [\" \".join(new_term)]\n",
    "        return new_terms\n",
    "    \n",
    "    # we need to decide whether to add to subj or obj.\n",
    "    second_text_pp = False\n",
    "    if modify_subject:\n",
    "        subject = text_split[0] if start_with_pp else text_split[1]\n",
    "        verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "        noun_pp = random.choice(list(noun_pool))\n",
    "        articles = 'a' if random.random() < 0.5 else 'the'\n",
    "        text_pp = [verb_pp_map[verb_pp], articles, noun_pp]\n",
    "        if articles == 'a':\n",
    "            pp_np_term =  [f\"{noun_pp} ( x _ 4 )\"]\n",
    "        elif articles == 'the':\n",
    "            pp_np_term = [f\"* {noun_pp} ( x _ 4 )\"]\n",
    "        pp_role_term = [f\"{subject} . nmod . {verb_pp} ( x _ 1 , x _ 4 )\"]\n",
    "        \n",
    "        if random.random() < 0.5:\n",
    "            second_text_pp = True\n",
    "            second_verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "            second_noun_pp = random.choice(list(noun_pool))\n",
    "            second_articles = 'a' if random.random() < 0.5 else 'the'\n",
    "            text_pp += [verb_pp_map[second_verb_pp], second_articles, second_noun_pp]\n",
    "            if second_articles == 'a':\n",
    "                pp_np_term += [f\"{second_noun_pp} ( x _ 7 )\"]\n",
    "            elif second_articles == 'the':\n",
    "                pp_np_term += [f\"* {second_noun_pp} ( x _ 7 )\"]\n",
    "            pp_role_term += [f\"{noun_pp} . nmod . {second_verb_pp} ( x _ 4 , x _ 7 )\"]\n",
    "            \n",
    "        def_terms = reindex_terms(def_terms, distance=6 if second_text_pp else 3)\n",
    "        rest_terms = reindex_terms(rest_terms, distance=6 if second_text_pp else 3)\n",
    "        \n",
    "        object_pos = 1\n",
    "        \n",
    "    else:\n",
    "        # locate the obj\n",
    "        assert \"x _ \" in object_entvar\n",
    "        object_pos = int(object_entvar.split()[-1])\n",
    "        _object = text_split[object_pos]\n",
    "        \n",
    "        verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "        noun_pp = random.choice(list(noun_pool))\n",
    "        articles = 'a' if random.random() < 0.5 else 'the'\n",
    "        text_pp = [verb_pp_map[verb_pp], articles, noun_pp]\n",
    "        if articles == 'a':\n",
    "            pp_np_term = [f\"{noun_pp} ( x _ {object_pos+3} )\"]\n",
    "        elif articles == 'the':\n",
    "            pp_np_term = [f\"* {noun_pp} ( x _ {object_pos+3} )\"]\n",
    "        pp_role_term = [f\"{_object} . nmod . {verb_pp} ( x _ {object_pos} , x _ {object_pos+3} )\"]\n",
    "        \n",
    "        if random.random() < 0.5:\n",
    "            second_text_pp = True\n",
    "            second_verb_pp = random.choice(list(verb_pp_map.keys()))\n",
    "            second_noun_pp = random.choice(list(noun_pool))\n",
    "            second_articles = 'a' if random.random() < 0.5 else 'the'\n",
    "            text_pp += [verb_pp_map[second_verb_pp], second_articles, second_noun_pp]\n",
    "            if second_articles == 'a':\n",
    "                pp_np_term += [f\"{second_noun_pp} ( x _ {object_pos+6} )\"]\n",
    "            elif second_articles == 'the':\n",
    "                pp_np_term += [f\"* {second_noun_pp} ( x _ {object_pos+6} )\"]\n",
    "            pp_role_term += [f\"{noun_pp} . nmod . {second_verb_pp} ( x _ {object_pos+3} , x _ {object_pos+6} )\"]\n",
    " \n",
    "        def_terms = reindex_terms(def_terms, object_pos, distance=6 if second_text_pp else 3)\n",
    "        rest_terms = reindex_terms(rest_terms, object_pos, distance=6 if second_text_pp else 3)\n",
    "    \n",
    "    for np_t in pp_np_term:\n",
    "        if np_t[0] == \"*\":\n",
    "            def_terms += [np_t]\n",
    "        else:\n",
    "            rest_terms += [np_t]\n",
    "    rest_terms += pp_role_term\n",
    "    text_split = text_split[:object_pos+1] + text_pp + text_split[object_pos+1:]\n",
    "    new_text = \" \".join(text_split)\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))   \n",
    "    rest_terms.sort(\n",
    "        key = lambda x: (\n",
    "            int(x.split()[-7]) if x.split()[-1] == \"nmod\" else int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), \n",
    "            int(x.split()[-3]) if x.split()[-1] == \"nmod\" else -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else new_text.split().index(x.split()[-2])\n",
    "        )\n",
    "    )  \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if second_text_pp:\n",
    "        return new_text, terms, \"subj_participle_verb_phrase_nested\" if modify_subject else \"obj_participle_verb_phrase_nested\"\n",
    "    else:\n",
    "        return new_text, terms, \"subj_participle_verb_phrase\" if modify_subject else \"obj_participle_verb_phrase\"\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51b9eebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "participial_prob = 0.15\n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "with open('./verb_pp_map.json', 'r') as openfile:\n",
    "    verb_pp_map = json.load(openfile)\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(\n",
    "    lambda x: collect_nv(*x), axis=1, result_type='expand'\n",
    ")\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(\n",
    "    lambda x: translate(*x), axis=1, result_type='expand'\n",
    ")\n",
    "\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a182c7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_postfix = \"participle_verb_easy\"\n",
    "train_df.to_csv(f'./cogs_{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./cogs_{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./cogs_{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./cogs_{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f6c030",
   "metadata": {},
   "source": [
    "### Set-based evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e82f2fdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv(\"./recogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./recogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./recogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./recogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36dd5a8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for lf in train_df[\"LF\"]:\n",
    "    if \"LAMBDA\" not in lf and translate_invariant_form_neoD(lf) == {}:\n",
    "        print(lf)\n",
    "for lf in gen_df[\"LF\"]:\n",
    "    if \"LAMBDA\" not in lf and translate_invariant_form_neoD(lf) == {}:\n",
    "        print(lf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "563ad239",
   "metadata": {},
   "source": [
    "### ReCOGS without coindexing change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c7dc48",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(seed)\n",
    "def translate_regular(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    return text, terms, _type\n",
    "        \n",
    "def translate(text, phi, _type):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, phi, _type\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    nmod_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                role_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            role_terms += [f\"{d['pred']} . {d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            assert \"x _\" in d['e1']\n",
    "            role_terms += [f\"{d['nppred']} . nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "    \n",
    "    if \"nmod\" not in terms.split():\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if random.random() >= proposing_prob:\n",
    "        return text, terms, _type\n",
    "    \n",
    "    if terms.split().count(\"nmod\") == 2:\n",
    "        upper_bound = 6\n",
    "        return_type = \"preposing_2\"\n",
    "    elif terms.split().count(\"nmod\") == 1:\n",
    "        upper_bound = 3\n",
    "        return_type = \"preposing_1\"\n",
    "    else:\n",
    "        assert False\n",
    "        \n",
    "    nmod = terms.split()[terms.split().index(\"nmod\")+2]\n",
    "    pre_phrase = text.split()[text.split().index(nmod)-2 : text.split().index(nmod)+upper_bound]\n",
    "    pre_phrase[0] = pre_phrase[0].capitalize()\n",
    "    pre_text = text.split()[:text.split().index(nmod)-2]\n",
    "    if pre_text[0] in [\"The\", \"A\"]:\n",
    "        pre_text[0] = pre_text[0].lower()\n",
    "    post_text = text.split()[text.split().index(nmod)+upper_bound:]\n",
    "    pre_text = pre_phrase + pre_text + post_text\n",
    "    pre_text = \" \".join(pre_text)\n",
    "\n",
    "    index_map = {}\n",
    "    idx = 0\n",
    "    for i in range(text.split().index(nmod)-2, text.split().index(nmod)+upper_bound):\n",
    "        index_map[f\"{i}\"] = f\"{idx}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)-2):\n",
    "        ii = 2+upper_bound+i\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"\n",
    "        idx += 1\n",
    "    for i in range(text.split().index(nmod)+upper_bound, len(text.split())):\n",
    "        ii = idx+(i-(text.split().index(nmod)+upper_bound))\n",
    "        index_map[f\"{i}\"] = f\"{ii}\"  \n",
    "            \n",
    "    # now handle LF\n",
    "    pre_terms = []\n",
    "    for t in terms.split():\n",
    "        if t.isnumeric():\n",
    "            pre_terms += [str(index_map[str(int(t))])]\n",
    "        else:\n",
    "            pre_terms += [t]\n",
    "    pre_terms = \" \".join(pre_terms)\n",
    "\n",
    "    pre_terms_def = pre_terms.split(\" ; \")[:-1]\n",
    "    pre_terms_def.sort(key = lambda x: int(x.split()[-2]))  \n",
    "    pre_terms_role = pre_terms.split(\" ; \")[-1].split(\" AND \")\n",
    "    pre_terms_role.sort(key = lambda x: (int(x.split()[-2]) if x.split()[-5] == \"(\" else int(x.split()[-6]) if x.split()[-2].isnumeric() else int(x.split()[-4]), -1 if x.split()[-5] == \"(\" else int(x.split()[-2]) if x.split()[-2].isnumeric() else pre_text.split().index(x.split()[-2])))  \n",
    "    pre_terms_role = \" AND \".join(pre_terms_role)\n",
    "    pre_terms = \" ; \".join(pre_terms_def + [pre_terms_role])\n",
    "\n",
    "    return pre_text, pre_terms, return_type\n",
    "\n",
    "def add_um(sentence):\n",
    "    words = sentence.split()\n",
    "    new_words = []\n",
    "    mapping = {}\n",
    "    offset = 0\n",
    "    for i, word in enumerate(words):\n",
    "        mapping[i] = len(new_words)\n",
    "        new_words.append(word)\n",
    "        if i > 0 and i < len(words) - 2 and random.random() > 0.5:\n",
    "            num_um = random.choice([1,2,3])\n",
    "            for j in range(num_um):\n",
    "                new_words.append(\"um\")\n",
    "    return \" \".join(new_words), mapping\n",
    "\n",
    "def sprinkle(text, phi, _type):\n",
    "    if \"preposition\" in _type:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    if random.random() >= sprinkle_prob:\n",
    "        return text, phi, _type\n",
    "    \n",
    "    um_text, token_mapping = add_um(text)\n",
    "    um_phi = []\n",
    "    for t in phi.split():\n",
    "        if t.isnumeric():\n",
    "            um_phi += [str(token_mapping[int(t)])]\n",
    "        else:\n",
    "            um_phi += [t]\n",
    "    um_phi = \" \".join(um_phi)\n",
    "    \n",
    "    return um_text, um_phi, f\"{_type}+sprinkle\"\n",
    "            \n",
    "train_df = pd.read_csv(\"./cogs/train.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "dev_df = pd.read_csv(\"./cogs/dev.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "test_df = pd.read_csv(\"./cogs/test.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "gen_df = pd.read_csv(\"./cogs/gen.tsv\", sep=\"\\t\", names=['sentence', 'LF', 'type'])\n",
    "\n",
    "proposing_prob = 0.05\n",
    "sprinkle_prob = 0.05\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "train_df[['sentence', 'LF', 'type']] = train_df[['sentence', 'LF', 'type']].apply(lambda x: sprinkle(*x), axis=1, result_type='expand')\n",
    "dev_df[['sentence', 'LF', 'type']] = dev_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF', 'type']] = test_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF', 'type']] = gen_df[['sentence', 'LF', 'type']].apply(lambda x: translate_regular(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be989297",
   "metadata": {},
   "outputs": [],
   "source": [
    "existing_digit_pool = set([])\n",
    "# loading target vocab to random sample our variable names\n",
    "for k, v in load_vocab(\"./model/tgt_vocab.txt\").items():\n",
    "    if k.isnumeric():\n",
    "        existing_digit_pool.add(k)\n",
    "existing_digit_pool = list(existing_digit_pool)\n",
    "\n",
    "def translate(text, phi):\n",
    "    \n",
    "    if len(phi.split()) == 1:\n",
    "        return text, f\"LAMBDA a . {phi} ( a )\"\n",
    "    elif \"LAMBDA\" in phi:\n",
    "        if len(phi.split()) == 7:\n",
    "            return text, phi\n",
    "        phi_split = phi.split(text)\n",
    "        cleaned_phi = []\n",
    "        for chunk in phi_split:\n",
    "            if \"LAMBDA\" in chunk:\n",
    "                cleaned_phi += [chunk.strip()]\n",
    "            else:\n",
    "                verb_args = chunk.strip(\" .\").split()[2]\n",
    "                cleaned_phi += [chunk.strip(\" .\")]\n",
    "        return text, \" \".join(cleaned_phi[:1] + [f\"{text} ( {verb_args} ) AND\"] + cleaned_phi[1:])\n",
    "    \n",
    "    # parse\n",
    "    text_split = text.split()\n",
    "    data = []    \n",
    "    conjs = re.split(r\"\\s*(?:AND|;)\\s*\", phi)\n",
    "    for conj in conjs: \n",
    "        if np_re.search(conj):\n",
    "            d = parse_np(conj)\n",
    "        elif pred_re.search(conj):\n",
    "            d = parse_pred(conj)\n",
    "            if \"x _\" not in d['entvar']:\n",
    "                d['entvar_name'] = d['entvar']\n",
    "                assert text_split.count(d['entvar']) == 1\n",
    "                name_idx = text_split.index(d['entvar'])\n",
    "                d['entvar'] = f\"x _ {name_idx}\"\n",
    "        elif mod_re.search(conj):\n",
    "            d = parse_mod(conj)\n",
    "        else:\n",
    "            raise Exception(f\"Conjunct could not be parsed: {conj}\")\n",
    "        data.append(d)\n",
    "    \n",
    "    # collect\n",
    "    def_terms = []\n",
    "    role_terms = []\n",
    "    for d in data:\n",
    "        if d['type'] == 'np':\n",
    "            if d['definiteness'] == '*':\n",
    "                def_terms += [f\"* {d['pred']} ( {d['entvar']} )\"]\n",
    "            else:\n",
    "                def_terms += [f\"{d['pred']} ( {d['entvar']} )\"]\n",
    "        if d['type'] == 'role':\n",
    "            if f\"{d['pred']} ( {d['eventvar']} )\" not in role_terms:\n",
    "                role_terms += [f\"{d['pred']} ( {d['eventvar']} )\"]\n",
    "            role_terms += [f\"{d['role']} ( {d['eventvar']} , {d['entvar']} )\"]\n",
    "            if \"entvar_name\" in d:\n",
    "                def_terms += [f\"{d['entvar_name']} ( {d['entvar']} )\"]\n",
    "        elif d['type'] == 'mod':\n",
    "            role_terms += [f\"nmod . {d['pred']} ( {d['e1']} , {d['e2']} )\"]\n",
    "            \n",
    "    # sort def_terms\n",
    "    def_terms = [*set(def_terms)]\n",
    "    def_terms.sort(key = lambda x: int(x.split()[-2]))    \n",
    "\n",
    "    rest_terms = role_terms\n",
    "    \n",
    "    # combine\n",
    "    def_terms = \" ; \".join(def_terms)\n",
    "    if def_terms == \"\":\n",
    "        terms = \" AND \".join(rest_terms)\n",
    "    elif \" AND \".join(rest_terms) == \"\":\n",
    "        terms = def_terms\n",
    "    else:\n",
    "        terms = def_terms + \" ; \" + \" AND \".join(rest_terms)\n",
    "\n",
    "    new_terms = []\n",
    "    for t in terms.split():\n",
    "        if t == \"_\" or t == \"x\":\n",
    "            continue\n",
    "        if t.isnumeric():\n",
    "            new_terms += [t]\n",
    "        else:\n",
    "            new_terms += [t]\n",
    "\n",
    "    terms = \" \".join(new_terms)\n",
    "    return text, terms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85c6e5bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampled_n = 1\n",
    "append_k = 3072\n",
    "\n",
    "train_df[['sentence', 'LF']] = train_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "dev_df[['sentence', 'LF']] = dev_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "test_df[['sentence', 'LF']] = test_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')\n",
    "gen_df[['sentence', 'LF']] = gen_df[['sentence', 'LF']].apply(lambda x: translate(*x), axis=1, result_type='expand')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc36f305",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reindex(LFs, initial_indexes):\n",
    "    new_LF_prefix = []\n",
    "    new_LF_body = []\n",
    "    \n",
    "    reindex_LFs = []\n",
    "    for i in range(len(LFs)):\n",
    "        if initial_indexes[i] != 0:\n",
    "            new_lf = []\n",
    "            for item in LFs[i].split():\n",
    "                if item.isnumeric():\n",
    "                    new_i = int(item) + initial_indexes[i]\n",
    "                    new_lf += [str(new_i)]\n",
    "                else:\n",
    "                    new_lf += [item]\n",
    "            new_lf = \" \".join(new_lf)\n",
    "        else:\n",
    "            new_lf = LFs[i]\n",
    "        reindex_LFs += [new_lf]\n",
    "        \n",
    "    new_LF_prefix = []\n",
    "    new_LF_body = []\n",
    "        \n",
    "    for i in range(len(reindex_LFs)):\n",
    "        new_LF_prefix.extend(reindex_LFs[i].split(\" ; \")[:-1])\n",
    "        for term in reindex_LFs[i].split(\" ; \")[-1].split(\" AND \"):\n",
    "            new_LF_body += [term]\n",
    "        \n",
    "    return \" ; \".join(new_LF_prefix) + \" ; \" + \" AND \".join(new_LF_body)\n",
    "\n",
    "start_indexes = [i*6 for i in range(append_k)]\n",
    "append_data = []\n",
    "\n",
    "for i in range(sampled_n):\n",
    "    sorted_train_df = train_df.sort_values(by=\"sentence\", key=lambda x: x.str.len())\n",
    "    for start_index in start_indexes:\n",
    "        conj_1 = sorted_train_df.iloc[-2-start_index].sentence\n",
    "        if conj_1.split()[0] in {'The', 'A'}:\n",
    "            conj_1_first = conj_1[0].lower()\n",
    "        else:\n",
    "            conj_1_first = conj_1[0]\n",
    "            \n",
    "        conj_2 = sorted_train_df.iloc[-3-start_index].sentence\n",
    "        if conj_2.split()[0] in {'The', 'A'}:\n",
    "            conj_2_first = conj_2[0].lower()\n",
    "        else:\n",
    "            conj_2_first = conj_2[0]\n",
    "            \n",
    "        append_data += [\n",
    "            [sorted_train_df.iloc[-1-start_index].sentence[:-1]+\\\n",
    "            conj_1_first+\\\n",
    "            sorted_train_df.iloc[-2-start_index].sentence[1:-1]+\\\n",
    "            conj_2_first+\\\n",
    "            sorted_train_df.iloc[-3-start_index].sentence[1:],\n",
    "            reindex(\n",
    "                [\n",
    "                    sorted_train_df.iloc[-1-start_index].LF,\n",
    "                    sorted_train_df.iloc[-2-start_index].LF,\n",
    "                    sorted_train_df.iloc[-3-start_index].LF\n",
    "                ],\n",
    "                [\n",
    "                    0,\n",
    "                    len(sorted_train_df.iloc[-1-start_index].sentence[:-1].strip().split()),\n",
    "                    len(sorted_train_df.iloc[-1-start_index].sentence[:-1].strip().split())+\n",
    "                    len(sorted_train_df.iloc[-2-start_index].sentence[:-1].strip().split())\n",
    "                ]\n",
    "            ),\n",
    "            'concat']\n",
    "        ]\n",
    "        \n",
    "append_data = pd.DataFrame(append_data, columns =['sentence', 'LF', 'type'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23efa7b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.concat([train_df, append_data])\n",
    "train_df = train_df.drop_duplicates()\n",
    "\n",
    "dataset_postfix = \"recogs_positional_index\"\n",
    "train_df.to_csv(f'./{dataset_postfix}/train.tsv', sep='\\t', index=False, header=False)\n",
    "dev_df.to_csv(f'./{dataset_postfix}/dev.tsv', sep='\\t', index=False, header=False)\n",
    "test_df.to_csv(f'./{dataset_postfix}/test.tsv', sep='\\t', index=False, header=False)\n",
    "gen_df.to_csv(f'./{dataset_postfix}/gen.tsv', sep='\\t', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefa5030",
   "metadata": {},
   "source": [
    "### Verification of paraphrased participle verb from ChatGPT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5891d346",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"en_core_web_sm\")\n",
    "f = open('./verb_pp_map.json')\n",
    "data = json.load(f)\n",
    "lemma_mapping = []\n",
    "hand_checks = {\"nurse\", \"bless\", \"ship\"}\n",
    "for k, v in data.items():\n",
    "    doc = nlp(v) # honestly, this is kind of waste, but whatever.\n",
    "    for token in doc:\n",
    "        if token.pos_ == \"VERB\":\n",
    "            lemma_mapping += [(k, token.lemma_)]\n",
    "        else:\n",
    "            if k not in hand_checks:\n",
    "                print(k, v)\n",
    "                assert False # hand check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "790ec7e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "for p in lemma_mapping:\n",
    "    if p[0] != p[1]:\n",
    "        print(p)\n",
    "        assert False"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}