scripts/generate_data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Generate data: Tatoeba\n",
    "This script generates clean and filtered data from data sources. The only data source currently is Tatoeba.\n",
    "\n",
    "Outputs are TSV files in the `data-generated/` folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%pip install pandas tqdm spacy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import csv\n",
    "import gc\n",
    "import re\n",
    "\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%sh\n",
    "mkdir data-generated"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Which sentences to keep by language.\n",
    "# If no sentence pairs are specified (below), all sentences for the given languages are kept. This does not guarantee that all sentences will have translations within the kept languages.\n",
    "# This variable maps 3-letter ISO language codes used by Tatoeba to 2-letter used by Taskbase.\n",
    "WANTED_LANGS = dict(\n",
    "    deu=\"DE\",\n",
    "    ukr=\"UK\",\n",
    ")\n",
    "\n",
    "# Optional. Whitelists of sentence pairs to keep.\n",
    "# This is useful if you are generating translation exercises and want to keep only sentences that are guaranteed to have a translation.\n",
    "# All languages appearing in these pair files must be present in WANTED_LANGS.\n",
    "# If this list is empty, all sentences having the languages in WANTED_LANGS will be included.\n",
    "# The files should be tab-separated with these columns: sentence 1 ID, sentence 1 text, sentence 2 ID, sentence 2 text.\n",
    "# Download Tatoeba sentence pair files from https://tatoeba.org/en/downloads\n",
    "SENTENCE_PAIR_FILES = [\n",
    "    \"./data-tatoeba/sentences_uk_de.tsv\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def build_tatoeba_sentences():\n",
    "    print(\"Tatoeba: building sentences...\")\n",
    "\n",
    "    # Get the master list of sentences.\n",
    "    df_sentences = pd.read_csv(\"./data-tatoeba/sentences_detailed.csv\", sep=\"\\t\",\n",
    "                               usecols=range(4), names=[\"id\", \"language\", \"text\", \"author\"],\n",
    "                               index_col=\"id\", quoting=csv.QUOTE_NONE)\n",
    "    df_sentences = df_sentences.dropna()\n",
    "\n",
    "    # Filter only wanted sentences\n",
    "    if len(SENTENCE_PAIR_FILES) > 0:\n",
    "        # If there are SENTENCE_PAIR_FILES, intersect those sentences with the master list. We need the intersection because the master list includes sentence metadata.\n",
    "        # Sentence pair files and the master list are updated at different times. Sometimes, sentences may be present in the pair files that aren't in the master list. If this happens as is important, make sure to re-download up-to-date datasets.\n",
    "        print(\"Filtering sentences by wanted pairs...\")\n",
    "        ids_to_keep = set()\n",
    "        for pair_file in SENTENCE_PAIR_FILES:\n",
    "            df_pairs = pd.read_csv(pair_file, sep=\"\\t\", names=[\"id1\", \"text1\", \"id2\", \"text2\"])\n",
    "            for id1, id2 in zip(df_pairs.id1, df_pairs.id2):\n",
    "                # Make sure both sentences are in the master list\n",
    "                if id1 in df_sentences.index and id2 in df_sentences.index:\n",
    "                    ids_to_keep.update([id1, id2])\n",
    "        df_sentences = df_sentences[df_sentences.index.isin(ids_to_keep)]\n",
    "    else:\n",
    "        # If no pair files are given, retain only the wanted languages.\n",
    "        print(\" Filtering sentences by language...\")\n",
    "        df_sentences = df_sentences[df_sentences.language.isin(WANTED_LANGS.keys())]\n",
    "\n",
    "    print(\"Replacing 'smart' quotes\")\n",
    "    re_quote = re.compile(\"[\\u201c\\u201d\\u201e\\u201f]\")\n",
    "    def mapper(text):\n",
    "        return re_quote.sub(\"\\\"\", text)\n",
    "    df_sentences.text = df_sentences.text.map(mapper)\n",
    "\n",
    "    # Map language codes\n",
    "    df_sentences.language = df_sentences.language.map(WANTED_LANGS)\n",
    "\n",
    "    print(f\"There are {len(df_sentences)} sentences\")\n",
    "\n",
    "    # Populate the \"translated_from\" field, for completeness.\n",
    "    # -1 means the base sentences was marked as null in the source dataset.\n",
    "    # -2 means the base sentence was missing in the source dataset.\n",
    "    # 0 means this is the root text.\n",
    "    print(\"Determining base sentences...\")\n",
    "    csv_base = pd.read_csv(\"./data-tatoeba/sentences_base.csv\", sep=\"\\t\",\n",
    "                           names=[\"id\", \"translated_from\"], index_col=\"id\",\n",
    "                           dtype={\"translated_from\": \"object\"})\n",
    "    s_translated_from = csv_base.translated_from.replace(\"\\\\N\", -1).astype(\"int64\")\n",
    "    df_sentences = df_sentences.join(s_translated_from, how=\"left\")\n",
    "    df_sentences[\"translated_from\"] = df_sentences.translated_from.fillna(-2)\n",
    "\n",
    "    # Count words\n",
    "    print(\"Counting words...\")\n",
    "    r_whitespace = re.compile(r\"\\s+\")\n",
    "    # r_letter = re.compile(\"[A-Za-z]\")\n",
    "    def get_word_count(string):\n",
    "        # TODO: How best to split words in different languages? Should non-word tokens be ignored?\n",
    "        return len([x for x in r_whitespace.split(string)])\n",
    "        # return len([x for x in r_whitespace.split(string) if r_letter.search(x) is not None])\n",
    "    s_word_count = df_sentences.text.map(get_word_count)\n",
    "    df_sentences[\"word_count\"] = s_word_count\n",
    "\n",
    "    # Write out\n",
    "    df_sentences.to_csv(\"data-generated/sentences.tsv\", sep=\"\\t\")\n",
    "    return df_sentences\n",
    "\n",
    "df_sentences = build_tatoeba_sentences()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Create translation relations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def build_translations(sentence_ids, dedup=True):\n",
    "    translations = []\n",
    "\n",
    "    def is_valid_link(s1, s2):\n",
    "        if dedup: return s1 in sentence_ids and s2 in sentence_ids and s1 < s2\n",
    "        else: return s1 in sentence_ids and s2 in sentence_ids\n",
    "\n",
    "    # Reading the links file manually is much faster than with Pandas\n",
    "    with open(\"./data-tatoeba/links.csv\", \"r\") as f:\n",
    "        for line in f:\n",
    "            s1, s2 = [int(x.strip()) for x in line.strip().split(\"\\t\")]\n",
    "            if is_valid_link(s1, s2): translations.append((s1, s2))\n",
    "\n",
    "    df_translations = pd.DataFrame(translations, columns=[\"s1\", \"s2\"])\n",
    "    df_translations.to_csv(\"data-generated/translations.tsv\", sep=\"\\t\", index=False)\n",
    "\n",
    "build_translations(set(df_sentences.index))\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Create vocabulary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def build_vocabulary():\n",
    "    import re\n",
    "    r_nonalpha = re.compile(r\"[^\\w\\-']\")\n",
    "    r_whitespace = re.compile(r\"\\s+\")\n",
    "\n",
    "    word_links = set()\n",
    "    n_word = 0\n",
    "    v2id = {}\n",
    "\n",
    "    for _, sentence in df_sentences.iterrows():\n",
    "        sentence_id = sentence.name\n",
    "\n",
    "        # Split words\n",
    "        words = [r_nonalpha.sub(\"\", x) for x in r_whitespace.split(sentence.text)]\n",
    "\n",
    "        for word in words:\n",
    "            v = (sentence.language, word)\n",
    "            if v not in v2id:\n",
    "                v2id[v] = n_word\n",
    "                n_word += 1\n",
    "            word_links.add((sentence_id, v2id[v]))\n",
    "\n",
    "    # Save sentence-vocabulary table\n",
    "    pd.DataFrame(word_links, columns=[\"sentence_id\", \"vocabulary_id\"])\\\n",
    "            .to_csv(\"data-generated/sentence-vocabulary.tsv\", sep=\"\\t\", index=False)\n",
    "\n",
    "    # Save vocabulary table\n",
    "    vocabs = v2id.keys()\n",
    "    pd.DataFrame([{\"language\": language, \"word\": word, \"length\": len(word)} for language, word in vocabs],\n",
    "                 index=v2id.values())\\\n",
    "            .rename_axis(\"id\")\\\n",
    "            .to_csv(\"data-generated/vocabulary.tsv\", sep=\"\\t\", index=True)\n",
    "\n",
    "build_vocabulary()\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Lemmatize\n",
    "\n",
    "**NOTE:** This part is yet unused in the exercise generations and can hence also be skipped.\n",
    "\n",
    "Make sure to include spaCy models for all languages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%sh\n",
    "python -m spacy download de_dep_news_trf\n",
    "python -m spacy download en_core_web_trf\n",
    "python -m spacy download fr_dep_news_trf\n",
    "python -m spacy download uk_core_news_trf\n",
    "\n",
    "python -m spacy download de_core_news_sm\n",
    "python -m spacy download en_core_web_sm\n",
    "python -m spacy download fr_core_news_sm\n",
    "python -m spacy download uk_core_news_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Build on the vocabulary table\n",
    "# Columns: id, language, word, length\n",
    "df_sentences = pd.read_csv(\"data-generated/sentences.tsv\", sep=\"\\t\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import spacy\n",
    "\n",
    "# Lightweight models.\n",
    "nlp_models = {\n",
    "    \"DE\": spacy.load(\"de_core_news_sm\"),\n",
    "    \"EN\": spacy.load(\"en_core_web_sm\"),\n",
    "    \"FR\": spacy.load(\"fr_core_news_sm\"),\n",
    "    \"UK\": spacy.load(\"uk_core_news_sm\"),\n",
    "}\n",
    "\n",
    "# Transformer models. Quite slow.\n",
    "# nlp_models = {\n",
    "#     \"DE\": spacy.load(\"de_dep_news_trf\"),\n",
    "#     \"EN\": spacy.load(\"en_core_web_trf\"),\n",
    "#     \"FR\": spacy.load(\"fr_dep_news_trf\"),\n",
    "#     \"UK\": spacy.load(\"uk_core_news_trf\"),\n",
    "# }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def build_lemmata():\n",
    "    links = set() # [(sentence_id, lemma_id)]\n",
    "    n_lemma = 0\n",
    "    lemma2id = {}\n",
    "\n",
    "    working_df = df_sentences#[df_sentences.language == \"EN\"].iloc[0:100]\n",
    "    for _, sentence in tqdm(working_df.iterrows(), total=len(working_df)):\n",
    "        tokens = nlp_models[sentence.language](sentence.text)\n",
    "\n",
    "        for token in tokens:\n",
    "            if not token.is_alpha: continue\n",
    "\n",
    "            lemma = (sentence.language, token.lemma_.lower())\n",
    "            if lemma not in lemma2id:\n",
    "                lemma2id[lemma] = n_lemma\n",
    "                n_lemma += 1\n",
    "            links.add((sentence.id, lemma2id[lemma]))\n",
    "\n",
    "    # TODO: Filter short or unwanted lemmata?\n",
    "\n",
    "    # Save sentence-lemma table\n",
    "    links_df = pd.DataFrame(links, columns=[\"sentence_id\", \"lemma_id\"])\n",
    "    links_df.to_csv(\"data-generated/sentence-lemma.tsv\", sep=\"\\t\", index=False)\n",
    "\n",
    "    # Save vocabulary table\n",
    "    lemmata = lemma2id.keys()\n",
    "    lemmata_df = pd.DataFrame([{\"language\": language, \"word\": word} for language, word in lemmata],\n",
    "                 index=lemma2id.values())\\\n",
    "            .rename_axis(\"id\")\n",
    "    lemmata_df.to_csv(\"data-generated/lemmata.tsv\", sep=\"\\t\", index=True)\n",
    "\n",
    "    return lemmata_df, links_df\n",
    "\n",
    "lemmata_df, links_df = build_lemmata()\n",
    "lemmata_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "del nlp_models\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "del df_sentences\n",
    "gc.collect()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}