extract_spanish_dtb.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from discourse_tree_utils import *\n",
    "from collections import defaultdict\n",
    "from string import digits\n",
    "import glob, random, json\n",
    "import pandas as pd\n",
    "pd.set_option('display.max_colwidth', -1)\n",
    "ISO='ISO-8859-1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = '/home/ffajri/Data/ES_DTB/ALL/*/*.rs3'\n",
    "files = glob.glob(path)\n",
    "random.shuffle(files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAIN = files[:200]\n",
    "DEV = files [200:230]\n",
    "TEST = files[230:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_text(graph, span):\n",
    "    start_index= int(span[3])\n",
    "    end_index = int(span[4])+1\n",
    "    return ' '.join([graph.node[a]['label'] for a in graph.tokens[start_index:end_index]])\n",
    "\n",
    "def subprocess(graph, spans):\n",
    "    edu1 = get_text(graph, spans[0])\n",
    "    nuc1 = spans[0][1].translate(None, digits)\n",
    "    rel = spans[0][2]\n",
    "    \n",
    "    if len(spans) == 2:\n",
    "        edu2 = get_text(graph, spans[1])\n",
    "    else:\n",
    "        edu2 = []\n",
    "        for span in spans[1:]:\n",
    "            edu2.append(get_text(graph, span))\n",
    "        edu2 = ' '.join(edu2)\n",
    "    nuc2 = spans[-1][1].translate(None, digits)\n",
    "    \n",
    "    return edu1, edu2, nuc1+nuc2, rel\n",
    "    \n",
    "def process(graph, spans):\n",
    "    span_dict = defaultdict(list)\n",
    "    for span in spans:\n",
    "        span_dict[span[0]].append(span)\n",
    "    for span in spans:\n",
    "        temp = sorted(span_dict[span[0]], key=lambda x: x[3])\n",
    "        span_dict[span[0]] = temp\n",
    "        \n",
    "    edus1 = []; edus2 = []; nucs = []; rels = []\n",
    "    for key in span_dict.keys():\n",
    "        edu1, edu2, nuc, rel = subprocess(graph, span_dict[key])\n",
    "        if edu1=='' or edu2=='':\n",
    "            continue\n",
    "        edus1.append(edu1)\n",
    "        edus2.append(edu2)\n",
    "        nucs.append(nuc)\n",
    "        rels.append(rel)\n",
    "    \n",
    "    df = pd.DataFrame()\n",
    "    df['edu1']=edus1; df['edu2']=edus2; df['nuclear']=nucs; df['relation']=rels\n",
    "    \n",
    "    edus = []\n",
    "    for edu in get_edus(graph):\n",
    "        edus.append(graph.node[edu]['rst:text'])\n",
    "    \n",
    "    return df, edus\n",
    "    \n",
    "def write_segment(fname, array):\n",
    "    final_data={}\n",
    "    for idx, data in enumerate(array):\n",
    "        final_data[idx]=data\n",
    "    json.dump(final_data, open(fname, 'w'))\n",
    "\n",
    "def compute_and_save(TARGET, fname):\n",
    "    final_df = pd.DataFrame(columns=['edu1', 'edu2', 'nuclear', 'relation'])\n",
    "    all_edus = []\n",
    "    for file in TARGET:\n",
    "        #print (file)\n",
    "        graph = RSTGraph(file, iso=ISO)\n",
    "        spans = get_rst_spans(graph)\n",
    "        df, edus = process(graph, spans)\n",
    "        final_df = final_df.append(df, ignore_index=True)\n",
    "        all_edus.append(edus)\n",
    "    \n",
    "    #save rst nuclear relation\n",
    "    final_df.to_csv('/home/ffajri/Workspace/WhatDiscourse/segment/data/data_es/'+fname+'.csv', index=False,  encoding=ISO)\n",
    "    #save segmentation\n",
    "    write_segment('/home/ffajri/Workspace/WhatDiscourse/segment/data/data_es/'+fname+'_edu.json', all_edus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#NUCLEARITY and RELATION PREDICTION and SEGMENTATION\n",
    "\n",
    "compute_and_save(TRAIN, 'train')\n",
    "compute_and_save(DEV, 'dev')\n",
    "compute_and_save(TEST, 'test')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}