dubbing-notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import os\n",
    "import numpy as np\n",
    "from proscript import *\n",
    "import sentencepiece as spm\n",
    "import matplotlib.mlab as mlab\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "import ast\n",
    "from collections import Counter\n",
    "from dataclasses import dataclass\n",
    "from random import shuffle\n",
    "from shutil import copyfile\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#LOAD HEROES CORPUS PATH DATA\n",
    "proscript_data_file = \"heroes_segment_paths_eval.txt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Auxiliary read-write functions\n",
    "\n",
    "def puncEstimate(punc):\n",
    "    if '.' in punc:\n",
    "        return '.'\n",
    "    elif ',' in punc:\n",
    "        return ','\n",
    "    elif '?' in punc:\n",
    "        return '?'\n",
    "    elif '!' in punc:\n",
    "        return '.'\n",
    "    elif ':' in punc:\n",
    "        return '.'\n",
    "    elif ';' in punc:\n",
    "        return ','\n",
    "    elif '-' in punc:\n",
    "        return '.'\n",
    "    else:\n",
    "        return ''\n",
    "\n",
    "def get_proscript_transcript(input_proscript, punctuation_as_token=False, reduce_punc=True):\n",
    "    input_tokens = []\n",
    "\n",
    "    for i, w in enumerate(input_proscript.word_list):\n",
    "        if punctuation_as_token:\n",
    "            if w.punctuation_before:\n",
    "                if reduce_punc:\n",
    "                    input_tokens.append(puncEstimate(w.punctuation_before))\n",
    "                else:\n",
    "                    input_tokens.append(w.punctuation_before)\n",
    "            input_tokens.append(w.word)\n",
    "            if w.punctuation_after:\n",
    "                if reduce_punc:\n",
    "                    input_tokens.append(puncEstimate(w.punctuation_after))\n",
    "                else:\n",
    "                    input_tokens.append(w.punctuation_after)\n",
    "        else:\n",
    "            input_tokens.append(w.punctuation_before + w.word + w.punctuation_after)\n",
    "    return ' '.join(input_tokens)\n",
    "\n",
    "#for reading old style (punctuation_before only)\n",
    "def get_proscript_transcript_2(input_proscript):\n",
    "    tokens = []\n",
    "    for i, w in enumerate(input_proscript.word_list):\n",
    "        if w.punctuation_before:\n",
    "            tokens[-1] += w.punctuation_before\n",
    "        tokens.append(w.word)\n",
    "    del tokens[-1]\n",
    "    return ' '.join(tokens)\n",
    "    \n",
    "    \n",
    "\n",
    "def read_text_file(file):\n",
    "\twith open(file,'r') as f:\n",
    "\t\treturn f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PARSE PROSCRIPT DATA FILE (File with paths to segment proscripts)\n",
    "heroes_data_all = []\n",
    "all_available_proscript_path = []  #input (EN) proscript paths\n",
    "with open(proscript_data_file, 'r') as f:\n",
    "    for l in f:\n",
    "        es_txt_path, es_csv_path, en_txt_path, en_csv_path = l.split()\n",
    "        all_available_proscript_path.append(en_csv_path)\n",
    "        heroes_data_all.append((es_txt_path, es_csv_path, en_txt_path, en_csv_path))   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SPLIT HEROES DATASET INTO DEVELOPMENT AND EVALUATION SETS\n",
    "heroes_data_dev = []\n",
    "heroes_data_dev_train = []\n",
    "heroes_data_dev_val = []\n",
    "\n",
    "heroes_data_eval = []\n",
    "\n",
    "heroes_data_shuffled = heroes_data_all.copy()\n",
    "\n",
    "shuffle(heroes_data_shuffled)\n",
    "\n",
    "dev_ratio = 0.5\n",
    "index = 0\n",
    "while index < len(heroes_data_shuffled) * dev_ratio:\n",
    "    heroes_data_dev.append(heroes_data_shuffled[index])\n",
    "    index += 1\n",
    "    \n",
    "while index < len(heroes_data_shuffled):\n",
    "    heroes_data_eval.append(heroes_data_shuffled[index])\n",
    "    index += 1    \n",
    "    \n",
    "#Split Dev set to training and validation\n",
    "validation_size = 200\n",
    "index = 0\n",
    "while index < validation_size:\n",
    "    heroes_data_dev_val.append(heroes_data_dev[index])\n",
    "    index += 1\n",
    "    \n",
    "while index < len(heroes_data_dev):\n",
    "    heroes_data_dev_train.append(heroes_data_dev[index])\n",
    "    index += 1  \n",
    "    \n",
    "print(\"Development: %i (train: %i, val: %i)\"%(len(heroes_data_dev), len(heroes_data_dev_train), len(heroes_data_dev_val)))\n",
    "print(\"Evaluation: %i \"%(len(heroes_data_eval)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PREPARE TEXT DATASET FROM HEROES SEGMENTS, PREPARE PATH FILES FOR SETS\n",
    "heroes_txt_en_path = 'onmt-heroes/heroes.en.txt'\n",
    "heroes_txt_es_path = 'onmt-heroes/heroes.es.txt'\n",
    "\n",
    "def data2textfiles(dataset, en_txt_output, es_txt_output, get_en_transcript_from='txt', get_es_transcript_from='txt'):\n",
    "    with open(en_txt_output, 'w') as f_en, open(es_txt_output, 'w') as f_es:\n",
    "        for es_txt_path, es_csv_path, en_txt_path, en_csv_path in dataset:\n",
    "            if get_en_transcript_from=='txt':\n",
    "                en_transcript = read_text_file(en_txt_path).lower()\n",
    "            elif get_en_transcript_from=='csv':\n",
    "                p_en = Proscript()\n",
    "                p_en.from_file(en_csv_path)\n",
    "                en_transcript = get_proscript_transcript_2(p_en)\n",
    "                \n",
    "                \n",
    "            if get_es_transcript_from=='txt':\n",
    "                es_transcript = read_text_file(es_txt_path).lower()\n",
    "            elif get_es_transcript_from=='csv':\n",
    "                p_es = Proscript()\n",
    "                p_es.from_file(es_csv_path)\n",
    "                es_transcript = get_proscript_transcript_2(p_es)\n",
    "                \n",
    "\n",
    "            f_es.write(es_transcript + '\\n')\n",
    "            f_en.write(en_transcript + '\\n')\n",
    "            \n",
    "def data2datafile(dataset, datafile_output):\n",
    "    with open(datafile_output, 'w') as f:\n",
    "        for es_txt_path, es_csv_path, en_txt_path, en_csv_path in dataset:\n",
    "            f.write(\"%s\\t%s\\t%s\\t%s\\n\"%(es_txt_path, es_csv_path, en_txt_path, en_csv_path))\n",
    "                \n",
    "# data2textfiles(heroes_data_all, 'onmt-heroes/heroes_all.en.txt', 'onmt-heroes/heroes_all.es.txt')\n",
    "# data2textfiles(heroes_data_eval, 'onmt-heroes/heroes_eval.en.txt', 'onmt-heroes/heroes_eval.es.txt')\n",
    "# data2textfiles(heroes_data_dev, 'onmt-heroes/heroes_dev.en.txt', 'onmt-heroes/heroes_dev.es.txt')\n",
    "# data2textfiles(heroes_data_dev_train, 'onmt-heroes/heroes_dev_train.en.txt', 'onmt-heroes/heroes_dev_train.es.txt')\n",
    "# data2textfiles(heroes_data_dev_val, 'onmt-heroes/heroes_dev_val.en.txt', 'onmt-heroes/heroes_dev_val.es.txt')\n",
    "\n",
    "# data2datafile(heroes_data_eval, 'heroes_segment_paths_eval.txt')\n",
    "# data2datafile(heroes_data_dev_train, 'heroes_segment_paths_dev_train.txt')\n",
    "# data2datafile(heroes_data_dev_val, 'heroes_segment_paths_dev_val.txt')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PREPARE PunkProse DATASET FOR PROSODIC PUNCTUATION MODEL ADAPTATION FOR ENGLISH\n",
    "#dev_train to train_samples\n",
    "punkProse_train_samples_dir = '/Users/alp/Documents/Corpora/punk_data/heroes_corpus_Vsugardub/train_samples'\n",
    "for _, _, _, en_csv_path in heroes_data_dev_train:\n",
    "    dst_file = os.path.join(punkProse_train_samples_dir, os.path.basename(en_csv_path))\n",
    "    copyfile(en_csv_path, dst_file)\n",
    "    \n",
    "#dev_val to dev_samples\n",
    "punkProse_dev_samples_dir = '/Users/alp/Documents/Corpora/punk_data/heroes_corpus_Vsugardub/dev_samples'\n",
    "for _, _, _, en_csv_path in heroes_data_dev_val:\n",
    "    dst_file = os.path.join(punkProse_dev_samples_dir, os.path.basename(en_csv_path))\n",
    "    copyfile(en_csv_path, dst_file)\n",
    "    \n",
    "#eval to test_samples\n",
    "punkProse_test_samples_dir = '/Users/alp/Documents/Corpora/punk_data/heroes_corpus_Vsugardub/test_samples'\n",
    "for _, _, _, en_csv_path in heroes_data_eval:\n",
    "    dst_file = os.path.join(punkProse_test_samples_dir, os.path.basename(en_csv_path))\n",
    "    copyfile(en_csv_path, dst_file)\n",
    "\n",
    "punkProse_groundtruth_samples_dir = '/Users/alp/Documents/Corpora/punk_data/heroes_corpus_Vsugardub/test_groundtruth'\n",
    "for _, _, _, en_csv_path in heroes_data_eval:\n",
    "    \n",
    "    dst_file = os.path.join(punkProse_groundtruth_samples_dir, os.path.splitext(os.path.basename(en_csv_path))[0] + '.txt')\n",
    "    input_proscript = Proscript()\n",
    "    input_proscript.from_file(en_csv_path)\n",
    "    \n",
    "    with open(dst_file, 'w') as f:\n",
    "        f.write(get_proscript_transcript(input_proscript, punctuation_as_token=True, reduce_punc=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PREPARE TEXT DATASET FROM punctuation recovered HEROES SEGMENTS, PREPARE PATH FILES FOR SETS\n",
    "\n",
    "def to_punkprosed(path):\n",
    "    return '/Users/alp/phdCloud/playground/punkHeroes/out_heroes_sugardub/csv/' + os.path.splitext(os.path.basename(path))[0] + '_punkProsed.csv'\n",
    "\n",
    "heroes_data_eval_punkProsed = [(a,b,c,to_punkprosed(d)) for a,b,c,d in heroes_data_eval]  \n",
    "data2textfiles(heroes_data_eval_punkProsed, 'onmt-heroes/heroes_eval_punkProsed.en.txt', \n",
    "               'onmt-heroes/heroes_eval_punkProsed.es.txt', \n",
    "               get_en_transcript_from='csv', \n",
    "              get_es_transcript_from='txt')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Functions for unit counting\n",
    "#Syllable counter. (Not very accurate for Spanish)\n",
    "def no_syl_en(word):\n",
    "    count = 0\n",
    "    vowels = 'aeiouyáóúíéüö'\n",
    "    word = word.lower().strip(\".:;?!\")\n",
    "    if word[0] in vowels:\n",
    "        count +=1\n",
    "    for index in range(1,len(word)):\n",
    "        if word[index] in vowels and word[index-1] not in vowels:\n",
    "            count +=1\n",
    "    if word.endswith('e'):\n",
    "        count -= 1\n",
    "    if word.endswith('le'):\n",
    "        count+=1\n",
    "    if count == 0:\n",
    "        count +=1\n",
    "    return count\n",
    "\n",
    "def no_syl_es(word):\n",
    "    count = 0\n",
    "    vowels = 'aeiouüö'\n",
    "    vowels_acento = 'áóúíé'\n",
    "    dipthongs = ['ai', 'au', 'eu', 'ei', 'oi', 'ou', 'ia', 'ie', 'io', 'iu', 'ua', 'ue', 'ui', 'uo', 'ió', 'ié']\n",
    "    word = word.lower().strip(\".:;?!\")\n",
    "    if word[0] in vowels + vowels_acento :\n",
    "        count +=1\n",
    "    for index in range(1,len(word)):\n",
    "        #print(word[index-1:index+1])\n",
    "        if word[index] in vowels + vowels_acento and word[index-1:index+1] not in dipthongs:\n",
    "            #print(word[index])\n",
    "            count +=1\n",
    "#     if 'oos' in word:\n",
    "#         count+=1\n",
    "    if count == 0:\n",
    "        count +=1\n",
    "    return count\n",
    "\n",
    "#Count number of words, syllables and characters in each sample\n",
    "def get_unit_counts(dataset):\n",
    "    heroes_unit_counts_en = []  # (Word, syllables, character)\n",
    "    heroes_unit_counts_es = []  # (Word, syllables, character)\n",
    "\n",
    "    for es_txt_path, es_csv_path, en_txt_path, en_csv_path in dataset:\n",
    "        #read corresponding proscripts\n",
    "        en_proscript = Proscript()\n",
    "        en_proscript.from_file(en_csv_path)\n",
    "\n",
    "        es_proscript = Proscript()\n",
    "        es_proscript.from_file(es_csv_path)\n",
    "\n",
    "        words_en = [w.word for w in en_proscript.word_list]\n",
    "        words_es = [w.word for w in es_proscript.word_list]\n",
    "\n",
    "        syl_en = [no_syl_en(w) for w in words_en]\n",
    "        syl_es = [no_syl_es(w) for w in words_es]\n",
    "\n",
    "        ch_en = []\n",
    "        for w in words_en:\n",
    "            ch_en += w\n",
    "\n",
    "        ch_es = []\n",
    "        for w in words_es:\n",
    "            ch_es += w\n",
    "\n",
    "        heroes_unit_counts_en.append((len(words_en), sum(syl_en), len(ch_en)))\n",
    "        heroes_unit_counts_es.append((len(words_es), sum(syl_es), len(ch_es))) \n",
    "\n",
    "    return heroes_unit_counts_en, heroes_unit_counts_es"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.9071633237822345\n",
      "10.249856733524355\n"
     ]
    }
   ],
   "source": [
    "#ANALYZE UNIT COUNTS IN SETS\n",
    "# #Get counts for all sets (takes long)\n",
    "counts_heroes_data_all_en, counts_heroes_data_all_es = get_unit_counts(heroes_data_all)\n",
    "# counts_heroes_data_eval_en, counts_heroes_data_eval_es = get_unit_counts(heroes_data_eval)\n",
    "# counts_heroes_data_dev_en, counts_heroes_data_dev_es = get_unit_counts(heroes_data_dev)\n",
    "# counts_heroes_data_dev_train_en, counts_heroes_data_dev_train_es = get_unit_counts(heroes_data_dev_train)\n",
    "# counts_heroes_data_dev_val_en, counts_heroes_data_dev_val_es = get_unit_counts(heroes_data_dev_val)\n",
    "\n",
    "#print average word counts\n",
    "print(np.mean([c[0] for c in counts_heroes_data_all_en]))\n",
    "# print(np.mean([c[0] for c in counts_heroes_data_eval_en]))\n",
    "# print(np.mean([c[0] for c in counts_heroes_data_dev_en]))\n",
    "# print(np.mean([c[0] for c in counts_heroes_data_dev_train_en]))\n",
    "# print(np.mean([c[0] for c in counts_heroes_data_dev_val_en]))\n",
    "\n",
    "#print average syllable counts\n",
    "print(np.mean([c[1] for c in counts_heroes_data_all_en]))\n",
    "# print(np.mean([c[1] for c in counts_heroes_data_eval_en]))\n",
    "# print(np.mean([c[1] for c in counts_heroes_data_dev_en]))\n",
    "# print(np.mean([c[1] for c in counts_heroes_data_dev_train_en]))\n",
    "# print(np.mean([c[1] for c in counts_heroes_data_dev_val_en]))\n",
    "\n",
    "#plot histogram of segment lengths w.r.t word count\n",
    "# to_plot = [counts_heroes_data_all_en, counts_heroes_data_eval_en, counts_heroes_data_dev_en, counts_heroes_data_dev_train_en, counts_heroes_data_dev_val_en]\n",
    "# plot_no = 1\n",
    "# unit_id = 0 #word\n",
    "# plt.figure(figsize=(8,20))\n",
    "# for counts in to_plot:\n",
    "#     plt.subplot(len(to_plot), 1, plot_no)\n",
    "#     x = [count[unit_id] for count in counts]\n",
    "#     # the histogram of the data\n",
    "#     plt.hist(x,density=1, bins=50) \n",
    "#     plt.axis([0, 40, 0, 0.2]) \n",
    "#     plt.xlabel('# Unit')\n",
    "#     plt.ylabel('Probability')\n",
    "#     plot_no += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Word\n",
      "7.9071633237822345\n",
      "6.905730659025788\n",
      "\n",
      "Syllable\n",
      "10.249856733524355\n",
      "12.80974212034384\n",
      "mean ratio 1.310834443649147\n",
      "std 0.3972435107732328\n",
      "\n",
      "Char\n",
      "31.702865329512893\n",
      "29.389684813753583\n"
     ]
    }
   ],
   "source": [
    "#Get unit count averages through samples\n",
    "print(\"Word\")\n",
    "print(np.mean([c[0] for c in counts_heroes_data_all_en]))\n",
    "print(np.mean([c[0] for c in counts_heroes_data_all_es]))\n",
    "print(\"\\nSyllable\")\n",
    "syl_mean_en = np.mean([c[1] for c in counts_heroes_data_all_en])\n",
    "syl_mean_es = np.mean([c[1] for c in counts_heroes_data_all_es])\n",
    "print(syl_mean_en)\n",
    "print(syl_mean_es)\n",
    "syllable_ratios = np.array([c[1] for c in counts_heroes_data_all_es]) / np.array([c[1] for c in counts_heroes_data_all_en])\n",
    "print('mean ratio', np.mean(syllable_ratios))\n",
    "print('std', np.std(syllable_ratios))\n",
    "\n",
    "print(\"\\nChar\")\n",
    "print(np.mean([c[2] for c in counts_heroes_data_all_en]))\n",
    "print(np.mean([c[2] for c in counts_heroes_data_all_es]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#LOAD SENTENCE PIECING MODEL \n",
    "sp = spm.SentencePieceProcessor()\n",
    "sp.Load(\"/Users/alp/Documents/Corpora/ted_en-es/sentencepieced_tedheroes/tedheroes.model\")\n",
    "\n",
    "#LOAD TRANSLATION OUTPUT\n",
    "#translation_data_file = '/Users/alp/phdCloud/playground/dubbing-w-pause/onmt-heroes/test/ted-opt-heroes-noreset_step_74970.pt_heroes_eval_punkProsed.translation.out'\n",
    "translation_data_file = '/Users/alp/phdCloud/playground/dubbing-w-pause/onmt-heroes/test/ted-opt-heroes-noreset_step_74970.pt_heroes.translation.out'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PARSE TRANSLATION INFORMATION [of all sentences] WITH ATTENTION WEIGHTS\n",
    "#Also does a proof check as some samples from the dataset have minor problems\n",
    "all_input_pieces = []\n",
    "all_pred_pieces = []\n",
    "all_attention_weights = []\n",
    "all_attention_binary = []\n",
    "all_proscript = []\n",
    "all_proscript_path = []\n",
    "\n",
    "with open(translation_data_file, 'r') as f:\n",
    "    sent_index = 0\n",
    "    attention_first_line = True\n",
    "    for l in f:\n",
    "        if l.isspace():\n",
    "            continue\n",
    "        elif l.startswith((\"GOLD\", \"PRED SCORE\")):\n",
    "            continue\n",
    "        elif l.startswith(\"SENT\"):\n",
    "            input_pieces = ast.literal_eval(l[l.find(':') + 2:])\n",
    "            continue\n",
    "        elif l.startswith(\"PRED\"):\n",
    "            pred_pieces = l[l.find(':') + 2:].split()\n",
    "            continue\n",
    "\n",
    "        if attention_first_line:\n",
    "            attention_first_line = False\n",
    "            num_input_pieces = len(input_pieces)\n",
    "            num_pred_pieces = len(pred_pieces)\n",
    "            sent_attention_weights = np.empty([0, num_input_pieces])\n",
    "            sent_attention_binary = np.empty([0, num_input_pieces])\n",
    "            continue\n",
    "            \n",
    "        #take the rest to the attention matrices until EOS is seen\n",
    "        try:\n",
    "            weights = [float(w) if not w.startswith('*') else float(w[1:]) for w in l.split()[1:]]\n",
    "            binary = [1 if w.startswith('*') else 0 for w in l.split()[1:]]\n",
    "        except:\n",
    "            print(l)\n",
    "            break\n",
    "        \n",
    "        if not l.split()[0] == \"</s>\":\n",
    "            sent_attention_weights = np.vstack([sent_attention_weights, weights])\n",
    "            sent_attention_binary = np.vstack([sent_attention_binary, binary])\n",
    "        else:\n",
    "            #read corresponding proscript\n",
    "            input_proscript = Proscript()\n",
    "            input_proscript.from_file(all_available_proscript_path[sent_index])\n",
    "\n",
    "            #make sure it matches with proscript info\n",
    "            proscript_transcript = get_proscript_transcript(input_proscript)\n",
    "            subtitle_transcript = sp.decode_pieces(input_pieces)\n",
    "            \n",
    "            if proscript_transcript == subtitle_transcript:\n",
    "                #store information \n",
    "                all_proscript_path.append(all_available_proscript_path[sent_index])\n",
    "                all_proscript.append(input_proscript)\n",
    "                all_input_pieces.append(input_pieces)\n",
    "                all_pred_pieces.append(pred_pieces)\n",
    "                all_attention_weights.append(sent_attention_weights)\n",
    "                all_attention_binary.append(sent_attention_binary)\n",
    "            else:\n",
    "                print(\"Problem at %s\"%all_available_proscript_path[sent_index])\n",
    "                print(subtitle_transcript)\n",
    "                print(proscript_transcript)\n",
    "                print(\"-----\")\n",
    "            attention_first_line = True\n",
    "            sent_index += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#EXTRACT PHRASE STRUCTURE FROM INPUT SENT PROSCRIPT\n",
    "\n",
    "@dataclass\n",
    "#struct to store phrase information\n",
    "class Phrase:\n",
    "    label: int\n",
    "    tokens: list\n",
    "    transcript: str\n",
    "    start_time: float\n",
    "    end_time: float\n",
    "    pause_after: float\n",
    "#Sample phrase initialize: p1 = Phrase(tokens=['a', 'b'], transcript=\"a b\", start_time=0.0)\n",
    "        \n",
    "#Get prosodic structure from proscript representation\n",
    "def get_input_structure(input_proscript, min_segment_pause_interval=0.0):\n",
    "    input_tokens = []\n",
    "    input_token_labels = []\n",
    "    curr_phrase_id = 0\n",
    "\n",
    "    phrase_structure = []\n",
    "    curr_phrase_tokens = []\n",
    "    curr_phrase_transcript = ''\n",
    "    get_start = True\n",
    "    for i, w in enumerate(input_proscript.word_list):\n",
    "        if get_start:\n",
    "            #first token in phrase\n",
    "            phrase_start_time = w.start_time\n",
    "            get_start = False\n",
    "            \n",
    "        token = w.punctuation_before + w.word + w.punctuation_after #THIS IS NOT REALLY TOKEN BUT TEXT ENCLOSED IN WHITESPACE\n",
    "        curr_phrase_transcript += token + ' '\n",
    "        input_tokens.append(token)\n",
    "        input_token_labels.append(curr_phrase_id)\n",
    "        curr_phrase_tokens.append(w.word)\n",
    "        \n",
    "        \n",
    "        #if there's a pause after the current word or if it is the last word, close the phrase\n",
    "        if w.pause_after > min_segment_pause_interval or i == len(input_proscript.word_list) - 1:\n",
    "            phrase = Phrase(label=curr_phrase_id,\n",
    "                            tokens=curr_phrase_tokens, \n",
    "                            transcript=curr_phrase_transcript.strip(), \n",
    "                            start_time=phrase_start_time, \n",
    "                            end_time=w.end_time, \n",
    "                            pause_after=w.pause_after)\n",
    "            \n",
    "            phrase_structure.append(phrase)\n",
    "            \n",
    "            #close phrase\n",
    "            curr_phrase_id += 1\n",
    "            curr_phrase_tokens = []\n",
    "            curr_phrase_transcript = ''\n",
    "            get_start = True\n",
    "    \n",
    "    return input_tokens, input_token_labels, phrase_structure\n",
    "    \n",
    "#labels piece sequence with phrase labels w.r.t input token labels\n",
    "#REFACTOR: this function can be merged to get_input_structure()\n",
    "def token_to_piece_labels(input_pieces, input_tokens, input_token_labels):\n",
    "    #label pieced tokens w.r.t proscript phrase labels\n",
    "    input_piece_labels = []\n",
    "    proscript_token_index = 0\n",
    "    piece_index = 0\n",
    "    curr_token_pieces = sp.EncodeAsPieces(input_tokens[proscript_token_index])\n",
    "    while piece_index < len(input_pieces):\n",
    "        try:\n",
    "            token_piece = curr_token_pieces.pop(0)\n",
    "        except:\n",
    "            proscript_token_index += 1\n",
    "            curr_token_pieces = sp.EncodeAsPieces(input_tokens[proscript_token_index])\n",
    "            continue\n",
    "\n",
    "        if input_pieces[piece_index] == token_piece:\n",
    "            input_piece_labels.append(input_token_labels[proscript_token_index])\n",
    "            piece_index += 1\n",
    "        else:\n",
    "            print(\"somethings not right in token alignment\")\n",
    "\n",
    "    return input_piece_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ALIGNING INPUT LABELS TO PREDICTION TOKENS - THE PROPER WAY#Naive USING BINARY ATTENTION MATRIX (kopuk olabilir)\n",
    "def map_piece_labels_naive(input_piece_labels, pred_pieces, sent_attention_binary):\n",
    "    pred_piece_labels = []\n",
    "    for piece_no, pred_piece in enumerate(pred_pieces):\n",
    "        matching_input_piece_no = np.argmax(sent_attention_binary[piece_no])\n",
    "        matching_input_piece = input_pieces[matching_input_piece_no]\n",
    "        piece_label = input_piece_labels[matching_input_piece_no]\n",
    "        pred_piece_labels.append(piece_label)\n",
    "        #print(pred_piece + \" - \"  + matching_input_piece + \" \" + str(piece_label))\n",
    "    \n",
    "    return pred_piece_labels\n",
    "\n",
    "#recursive function to extend sequence possibilities\n",
    "def extend_possibilities(poss_label_seqs, labels, extend_by):\n",
    "    new_poss_label_seqs = []\n",
    "    for label_seq in poss_label_seqs:\n",
    "        #print('label_seq', label_seq)\n",
    "        for label in labels:\n",
    "            #print('label', label)\n",
    "            if label == label_seq[-1] or label == label_seq[-1] + 1:\n",
    "                new_seq = label_seq + [label]\n",
    "                #print('new_seq', new_seq)\n",
    "                new_poss_label_seqs.append(new_seq)\n",
    "    #print(\"...\")\n",
    "    if extend_by == 1:\n",
    "        return new_poss_label_seqs\n",
    "    else:\n",
    "        extended_poss_label_seqs = []\n",
    "        for seq in new_poss_label_seqs:\n",
    "            #print('seq', seq)\n",
    "            new_ones = extend_possibilities([seq], labels, extend_by - 1)\n",
    "            #print(new_ones)\n",
    "            extended_poss_label_seqs.extend(new_ones)\n",
    "        return extended_poss_label_seqs\n",
    "\n",
    "#check if pieces of same word have conflicting labels\n",
    "def check_splitlabels(piece_seq, piece_label_seq):\n",
    "    for i, piece in enumerate(piece_seq):\n",
    "        if not piece.startswith('▁'):\n",
    "            if not piece_label_seq[i - 1] == piece_label_seq[i]:\n",
    "                return False\n",
    "    return True\n",
    "\n",
    "def label_mask(label_seq, hot_label):\n",
    "    masked = []\n",
    "    for l in label_seq:\n",
    "        if l == hot_label:\n",
    "            masked.append(1)\n",
    "        else:\n",
    "            masked.append(0)\n",
    "    return masked  \n",
    "    \n",
    "def get_pred_labels(input_pieces, input_piece_labels, pred_pieces, sent_attention_weights):\n",
    "    labels = np.unique(input_piece_labels)\n",
    "    \n",
    "    all_possible_pred_label_seqs = [[0]]\n",
    "    all_possible_pred_label_seqs = extend_possibilities(all_possible_pred_label_seqs, labels, len(pred_pieces) - 1)\n",
    "    \n",
    "    #filter out seqs where input seqs are not fully represented and if there's any multilabeled words\n",
    "    possible_pred_label_seqs = [seq for seq in all_possible_pred_label_seqs if labels[-1] in seq and check_splitlabels(pred_pieces, seq)]\n",
    "    \n",
    "    #get scores of each possible sequence\n",
    "    possible_seq_scores = [1] * len(possible_pred_label_seqs)\n",
    "    for i_seq, poss_label_seq in enumerate(possible_pred_label_seqs):\n",
    "        #print('seq', poss_label_seq)\n",
    "        for pred_piece_index, label in enumerate(poss_label_seq):\n",
    "            #print('i:%i - l:%i'%(pred_piece_index,label))\n",
    "            masked_weights = sent_attention_weights[pred_piece_index] * label_mask(input_piece_labels, label)\n",
    "            weights_sum = np.sum(masked_weights)\n",
    "            #print('weights', masked_weights)\n",
    "            #print('weights_sum', weights_sum)\n",
    "            possible_seq_scores[i_seq] *= weights_sum\n",
    "        #print(\"--> seq score:\", possible_seq_scores[i_seq])    \n",
    "        #print('---')\n",
    "        \n",
    "    try:    \n",
    "        best_possible_sequence = possible_pred_label_seqs[np.argmax(possible_seq_scores)]\n",
    "    except:\n",
    "        best_possible_sequence = [0] * len(pred_pieces)\n",
    "        \n",
    "    return best_possible_sequence, possible_pred_label_seqs, possible_seq_scores\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#TRANSLATE PREDICTION PIECE LABELS TO PREDICTION TOKEN LABELS\n",
    "def most_common(lst):\n",
    "    data = Counter(lst)\n",
    "    return max(lst, key=data.get)\n",
    "\n",
    "def piece_to_token(pred_pieces, pred_piece_labels):\n",
    "    tokens = []\n",
    "    token_labels = []\n",
    "    token_label_votes = []\n",
    "    for piece_index, piece in enumerate(pred_pieces):\n",
    "        if piece.startswith('▁'):\n",
    "            tokens.append(piece[1:])\n",
    "            token_label_votes.append([pred_piece_labels[piece_index]])\n",
    "        else:\n",
    "            tokens[-1] += piece\n",
    "            token_label_votes[-1].append(pred_piece_labels[piece_index])\n",
    "    \n",
    "    token_labels = [most_common(token_votes) for token_votes in token_label_votes]\n",
    "    \n",
    "    return tokens, token_labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['vamos,', 'vamos.']\n",
      "['vamos', 'vamos']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Phrase(label=0, tokens=['vamos', 'vamos'], transcript='vamos vamos', start_time=0.0, end_time=9.73, pause_after=0.0)]"
      ]
     },
     "execution_count": 77,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred_phrase_structure, problem_matching = get_output_phrase_structure(pred_tokens, pred_token_labels, input_phrase_structure)\n",
    "pred_phrase_structure\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "#EXTRACT OUTPUT PHRASE STRUCTURE W.R.T INPUT PHRASE STRUCTURE\n",
    "def get_output_phrase_structure(output_tokens, output_token_labels, input_structure):\n",
    "    output_structure = []\n",
    "    curr_phrase_id = 0\n",
    "    curr_phrase_tokens = []\n",
    "    no_match = False\n",
    "    \n",
    "    #Take care of the case that output phrases cannot match all input phrases\n",
    "    if len(input_structure) > max(output_token_labels) + 1:\n",
    "        no_match = True\n",
    "        #make only one phrase\n",
    "        curr_phrase_tokens = output_tokens\n",
    "        phrase_tokens = [strip_punctuation(tok) for tok in curr_phrase_tokens]\n",
    "        print(curr_phrase_tokens)\n",
    "        print(phrase_tokens)\n",
    "        phrase = Phrase(label=0,\n",
    "                        tokens=phrase_tokens, \n",
    "                        transcript=' '.join(phrase_tokens), \n",
    "                        start_time=input_structure[0].start_time, \n",
    "                        end_time=input_structure[-1].end_time, \n",
    "                        pause_after=0.0)\n",
    "        output_structure.append(phrase)\n",
    "    else:\n",
    "        for i, (tok, l) in enumerate(zip(output_tokens, output_token_labels)):\n",
    "            #print(tok)\n",
    "            t = strip_punctuation(tok) \n",
    "            if not l == curr_phrase_id:\n",
    "                #close phrase\n",
    "                phrase = Phrase(label=curr_phrase_id,\n",
    "                                tokens=curr_phrase_tokens, \n",
    "                                transcript=' '.join(curr_phrase_tokens), \n",
    "                                start_time=input_structure[curr_phrase_id].start_time, \n",
    "                                end_time=input_structure[curr_phrase_id].end_time, \n",
    "                                pause_after=input_structure[curr_phrase_id].pause_after)\n",
    "\n",
    "                output_structure.append(phrase)\n",
    "                curr_phrase_id += 1\n",
    "                if t:\n",
    "                    curr_phrase_tokens = [t]\n",
    "            else:\n",
    "                if t:\n",
    "                    curr_phrase_tokens.append(t)\n",
    "        else:\n",
    "            #get the final phrase\n",
    "            phrase = Phrase(label=curr_phrase_id,\n",
    "                                tokens=curr_phrase_tokens, \n",
    "                                transcript=' '.join(curr_phrase_tokens), \n",
    "                                start_time=input_structure[curr_phrase_id].start_time, \n",
    "                                end_time=input_structure[curr_phrase_id].end_time, \n",
    "                                pause_after=input_structure[curr_phrase_id].pause_after)\n",
    "\n",
    "            output_structure.append(phrase)\n",
    "            \n",
    "    return output_structure, no_match\n",
    "\n",
    "def strictly_increasing(L):\n",
    "    return all(x<=y for x, y in zip(L, L[1:]))\n",
    "\n",
    "#strips punctuation from beginning and end of token\n",
    "def strip_punctuation(token_string):\n",
    "    word_being_processed = token_string\n",
    "    punc_after = \"\"\n",
    "    punc_before = \"\"\n",
    "    \n",
    "    if not re.search(r\"\\w\", word_being_processed):\n",
    "        return ''\n",
    "    \n",
    "    if re.search(r\"^\\W\", word_being_processed):\n",
    "        punc = word_being_processed[:re.search(r\"\\w\", word_being_processed).start()]\n",
    "        punc_before += punc\n",
    "        word_being_processed = word_being_processed[re.search(r\"\\w\", word_being_processed).start():]\n",
    "\n",
    "    #check end again (issue with quotations)\n",
    "    word_reversed = word_being_processed[::-1]\n",
    "    if re.search(r\"^\\W\",word_reversed):\n",
    "        punc = word_reversed[:re.search(r\"\\w\", word_reversed).start()][::-1]\n",
    "        punc_after = punc + punc_after\n",
    "        word_being_processed = word_reversed[re.search(r\"\\w\", word_reversed).start():][::-1]\n",
    "\n",
    "    return word_being_processed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Show attention matrix\n",
    "#Sample call: show_attention(input_pieces, pred_pieces, sent_attention_weights)\n",
    "color_palette = ['blue', 'green', 'red', 'brown', 'magenta', 'black', 'cyan']\n",
    "\n",
    "def show_attention(input_words, output_words, attentions, input_word_labels = None, pred_word_labels = None):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure(figsize=(12, 12))\n",
    "    \n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions, cmap='bone_r')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    ax.set_xticklabels([''] + input_words, rotation=90, fontsize=15)\n",
    "    ax.set_yticklabels([''] + output_words, fontsize=15)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    \n",
    "    #return plt.gca().get_xticklabels(), plt.gca().get_yticklabels()\n",
    "    \n",
    "    #Show labels of words as colors\n",
    "    if input_word_labels:\n",
    "        for i, word_tick in enumerate(plt.gca().get_xticklabels()[1:-1]):\n",
    "            word_tick.set_color(color_palette[input_word_labels[i] % len(color_palette)]) \n",
    "    \n",
    "    if pred_word_labels:\n",
    "        for i, word_tick in enumerate(plt.gca().get_yticklabels()[1:-1]):\n",
    "            word_tick.set_color(color_palette[pred_word_labels[i] % len(color_palette)]) \n",
    "        \n",
    "    plt.show()\n",
    "    plt.close()\n",
    "\n",
    "#sample call: show_attention(input_pieces, pred_pieces, sent_attention_weights, input_piece_labels, pred_piece_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "#CALCULATE AVERAGE NUMBER OF SYLLABLES SIMILARITY PENALTY\n",
    "\n",
    "syl_std = 0.38\n",
    "syl_mean = 1.30\n",
    "\n",
    "def get_syllable_similarity_penalty(utt_en, utt_es):\n",
    "    total_factor = 0\n",
    "    \n",
    "    if not len(utt_en) == len(utt_es):\n",
    "        en_tokens = []\n",
    "        for phrase in utt_en:\n",
    "            en_tokens.extend(phrase.tokens)\n",
    "            \n",
    "        es_tokens = []\n",
    "        for phrase in utt_es:\n",
    "            es_tokens.extend(phrase.tokens)\n",
    "            \n",
    "        syl_en = sum([no_syl_en(word) for word in en_tokens])\n",
    "        syl_es = sum([no_syl_es(word) for word in es_tokens])\n",
    "        \n",
    "        #print(phrase_en.transcript)\n",
    "        #print(phrase_es.transcript)\n",
    "        ratio = syl_es/syl_en\n",
    "        factor = ratio\n",
    "        \n",
    "        #print(\"en:%i, es:%i - %f, f:%f\"%(syl_en, syl_es, ratio, factor))\n",
    "        #print(\"...\")\n",
    "        total_factor += factor\n",
    "    else:    \n",
    "        for phrase_en, phrase_es in zip(utt_en, utt_es):\n",
    "            syl_en = sum([no_syl_en(word) for word in phrase_en.tokens])\n",
    "            syl_es = sum([no_syl_es(word) for word in phrase_es.tokens])\n",
    "\n",
    "            #print(phrase_en.transcript)\n",
    "            #print(phrase_es.transcript)\n",
    "            ratio = syl_es/syl_en\n",
    "            factor = ratio\n",
    "\n",
    "            #print(\"en:%i, es:%i - %f, f:%f\"%(syl_en, syl_es, ratio, factor))\n",
    "            #print(\"...\")\n",
    "            total_factor += factor\n",
    "    \n",
    "    avg_factor = total_factor / len(utt_en)\n",
    "    \n",
    "    normalized = (avg_factor - syl_mean) / syl_std\n",
    "    \n",
    "    penalty = abs(normalized)\n",
    "        \n",
    "    return penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "at 100\n",
      "Couldn't match all phrases in 111\n",
      "Couldn't match all phrases in 134\n",
      "at 200\n",
      "at 300\n",
      "Couldn't match all phrases in 367\n",
      "at 400\n",
      "at 500\n",
      "at 600\n",
      "Couldn't match all phrases in 643\n",
      "Couldn't match all phrases in 683\n",
      "at 700\n",
      "at 800\n",
      "at 900\n",
      "Couldn't match all phrases in 980\n",
      "at 1000\n",
      "at 1100\n",
      "at 1200\n",
      "Couldn't match all phrases in 1242\n",
      "at 1300\n",
      "Couldn't match all phrases in 1310\n",
      "at 1400\n",
      "Couldn't match all phrases in 1420\n",
      "at 1500\n",
      "at 1600\n",
      "at 1700\n",
      "Couldn't match all phrases in 1732\n",
      "at 1800\n",
      "Couldn't match all phrases in 1897\n",
      "at 1900\n",
      "Couldn't match all phrases in 1957\n",
      "at 2000\n",
      "at 2100\n",
      "at 2200\n",
      "at 2300\n",
      "at 2400\n",
      "Couldn't match all phrases in 2479\n",
      "at 2500\n",
      "at 2600\n",
      "at 2700\n",
      "Couldn't match all phrases in 2738\n",
      "at 2800\n",
      "at 2900\n",
      "Couldn't match all phrases in 2917\n",
      "at 3000\n",
      "at 3100\n",
      "at 3200\n",
      "at 3300\n",
      "Couldn't match all phrases in 3304\n",
      "at 3400\n",
      "Processed 3490. 0 problematic\n"
     ]
    }
   ],
   "source": [
    "#BATCH PROCESS INPUT SEGMENTS->OUTPUT SEGMENTS\n",
    "all_input_tokens = []\n",
    "all_input_token_labels = []\n",
    "all_input_piece_labels = []\n",
    "all_pred_piece_labels = []\n",
    "all_pred_tokens = []\n",
    "all_pred_token_labels = []\n",
    "all_input_phrase_structure = []\n",
    "all_pred_phrase_structure = []\n",
    "all_syllable_similarity_penalty = []\n",
    "\n",
    "MIN_SEGMENT_PAUSE_INTERVAL = 0.25\n",
    "OUTPUT_DIR = 'samples/batchpro'\n",
    "\n",
    "severe_problematic= 0\n",
    "count = 0\n",
    "\n",
    "for segment_no, (input_pieces, pred_pieces, input_proscript, sent_attention_binary, sent_attention_weights, proscript_path) in enumerate(zip(all_input_pieces, all_pred_pieces, all_proscript, all_attention_binary, all_attention_weights, all_proscript_path)):\n",
    "    try:\n",
    "        input_tokens, input_token_labels, input_phrase_structure = get_input_structure(input_proscript, MIN_SEGMENT_PAUSE_INTERVAL)\n",
    "        input_piece_labels = token_to_piece_labels(input_pieces, input_tokens, input_token_labels)\n",
    "        pred_piece_labels, _, _ = get_pred_labels(input_pieces, input_piece_labels, pred_pieces, sent_attention_weights)\n",
    "        pred_tokens, pred_token_labels = piece_to_token(pred_pieces, pred_piece_labels)\n",
    "        pred_phrase_structure, problem_matching = get_output_phrase_structure(pred_tokens, pred_token_labels, input_phrase_structure)\n",
    "\n",
    "        syllable_similarity_penalty = get_syllable_similarity_penalty(input_phrase_structure, pred_phrase_structure)\n",
    "    except:\n",
    "        input_tokens, input_token_labels, input_piece_labels, pred_piece_labels, pred_tokens, pred_token_labels = None, None, None, None, None, None\n",
    "        print(\"Somethings wrong processing %i\"%segment_no)\n",
    "        severe_problematic += 1\n",
    "        \n",
    "    if problem_matching:\n",
    "        print(\"Couldn't match all phrases in %i\"%segment_no)\n",
    "        \n",
    "    all_input_tokens.append(input_tokens)\n",
    "    all_input_token_labels.append(input_token_labels)\n",
    "    all_input_piece_labels.append(input_piece_labels)\n",
    "    all_input_phrase_structure.append(input_phrase_structure)\n",
    "    \n",
    "    all_pred_piece_labels.append(pred_piece_labels)\n",
    "    all_pred_tokens.append(pred_tokens)\n",
    "    all_pred_token_labels.append(pred_token_labels)\n",
    "    all_pred_phrase_structure.append(pred_phrase_structure)\n",
    "    \n",
    "    all_syllable_similarity_penalty.append(syllable_similarity_penalty)\n",
    "    \n",
    "    #output files for synthesis\n",
    "    #1. plain translation string\n",
    "    sample_id = ''.join(os.path.splitext(os.path.basename(proscript_path))[0].split(\"_eng_aligned\")).replace('eng', '').replace('heroes_', '')\n",
    "    translation_output_file = os.path.join(OUTPUT_DIR, sample_id + '.translation.transcript.txt')\n",
    "    full_translation = ' '.join(pred_tokens)\n",
    "    with open (translation_output_file, 'w') as f_txt:\n",
    "        f_txt.write(full_translation)\n",
    "\n",
    "    #2. phrase structure\n",
    "    phrase_output_file = os.path.join(OUTPUT_DIR, sample_id + '.translation.structure.txt')\n",
    "    with open(phrase_output_file, 'w') as f:\n",
    "        for st in pred_phrase_structure:\n",
    "            f.write(\"%f\\t%f\\t%f\\t%s\\n\"%(st.start_time, st.end_time, st.pause_after, st.tokens))\n",
    "\n",
    "    count += 1\n",
    "    if count % 100 == 0:    \n",
    "        print('at', count)\n",
    "    \n",
    "print(\"Processed %i. %i problematic\"%(count, severe_problematic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "111 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_2/spa-eng/segments_eng/heroes_s2_2_eng_aligned_eng0151.csv\n",
      "['¿qué?']\n",
      "['qué']\n",
      "syllable_similarity_penalty 2.5438596491228074\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['so', 'lizards']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁¿ qué ?']\n",
      "\n",
      "Segmented output tokens\n",
      "['¿qué?']\n",
      "\n",
      "Matched input and pred segments\n",
      "so lizards\n",
      "¿qué?\n",
      "\n",
      ">so\n",
      "<¿qué?\n",
      "\n",
      "134 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_5/spa-eng/segments_eng/heroes_s2_5_eng_aligned_eng0130.csv\n",
      "['¿dónde', 'aprendías?']\n",
      "['dónde', 'aprendías']\n",
      "syllable_similarity_penalty 2.669172932330827\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['well', 'where', 'did you learn how to']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁¿ dónde ▁aprend ías ?']\n",
      "\n",
      "Segmented output tokens\n",
      "['¿dónde aprendías?']\n",
      "\n",
      "Matched input and pred segments\n",
      "well where did you learn how to\n",
      "¿dónde aprendías?\n",
      "\n",
      ">well\n",
      "<¿dónde aprendías?\n",
      "\n",
      "367 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_6/spa-eng/segments_eng/heroes_s2_6_eng_aligned_eng0333.csv\n",
      "['mira.']\n",
      "['mira']\n",
      "syllable_similarity_penalty 2.763157894736842\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['so', 'look around']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁mira .']\n",
      "\n",
      "Segmented output tokens\n",
      "['mira.']\n",
      "\n",
      "Matched input and pred segments\n",
      "so look around\n",
      "mira.\n",
      "\n",
      ">so\n",
      "<mira.\n",
      "\n",
      "643 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_9/spa-eng/segments_eng/heroes_s3_9_eng_aligned_eng0149.csv\n",
      "['adelante,', 'matarme.']\n",
      "['adelante', 'matarme']\n",
      "syllable_similarity_penalty 2.192982456140351\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['go ahead', 'kill', 'me']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁adelante , ▁mat arme .']\n",
      "\n",
      "Segmented output tokens\n",
      "['adelante, matarme.']\n",
      "\n",
      "Matched input and pred segments\n",
      "go ahead kill me\n",
      "adelante, matarme.\n",
      "\n",
      ">go ahead\n",
      "<adelante, matarme.\n",
      "\n",
      "683 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_11/spa-eng/segments_eng/heroes_s3_11_eng_aligned_eng0066.csv\n",
      "['¡no!', '¡no!']\n",
      "['no', 'no']\n",
      "syllable_similarity_penalty 3.325358851674641\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['no note', 'no goodbye', 'just gone', 'again', 'you okay']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ no ! ▁ ¡ no !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡no! ¡no!']\n",
      "\n",
      "Matched input and pred segments\n",
      "no note no goodbye just gone again you okay\n",
      "¡no! ¡no!\n",
      "\n",
      ">no note\n",
      "<¡no! ¡no!\n",
      "\n",
      "980 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_7/spa-eng/segments_eng/heroes_s2_7_eng_aligned_eng0355.csv\n",
      "['¡tela!']\n",
      "['tela']\n",
      "syllable_similarity_penalty 2.763157894736842\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['shut', 'up lyle']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ t ela !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡tela!']\n",
      "\n",
      "Matched input and pred segments\n",
      "shut up lyle\n",
      "¡tela!\n",
      "\n",
      ">shut\n",
      "<¡tela!\n",
      "\n",
      "1242 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_11/spa-eng/segments_eng/heroes_s2_11_eng_aligned_eng0295.csv\n",
      "['¡vamos!']\n",
      "['vamos']\n",
      "syllable_similarity_penalty 2.8947368421052633\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['go', 'get the virus']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ vamos !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡vamos!']\n",
      "\n",
      "Matched input and pred segments\n",
      "go get the virus\n",
      "¡vamos!\n",
      "\n",
      ">go\n",
      "<¡vamos!\n",
      "\n",
      "1310 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_4/spa-eng/segments_eng/heroes_s2_4_eng_aligned_eng0419.csv\n",
      "['vale.']\n",
      "['vale']\n",
      "syllable_similarity_penalty 2.1052631578947367\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['ok', 'ok']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁vale .']\n",
      "\n",
      "Segmented output tokens\n",
      "['vale.']\n",
      "\n",
      "Matched input and pred segments\n",
      "ok ok\n",
      "vale.\n",
      "\n",
      ">ok\n",
      "<vale.\n",
      "\n",
      "1420 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_8/spa-eng/segments_eng/heroes_s3_8_eng_aligned_eng0182.csv\n",
      "['vamos.']\n",
      "['vamos']\n",
      "syllable_similarity_penalty 2.1052631578947367\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['come', 'on']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁vamos .']\n",
      "\n",
      "Segmented output tokens\n",
      "['vamos.']\n",
      "\n",
      "Matched input and pred segments\n",
      "come on\n",
      "vamos.\n",
      "\n",
      ">come\n",
      "<vamos.\n",
      "\n",
      "1732 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_9/spa-eng/segments_eng/heroes_s3_9_eng_aligned_eng0061.csv\n",
      "['vamos,', 'vamos.']\n",
      "['vamos', 'vamos']\n",
      "syllable_similarity_penalty 2.919799498746867\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['go', 'go', 'we should keep moving']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁vamos , ▁vamos .']\n",
      "\n",
      "Segmented output tokens\n",
      "['vamos, vamos.']\n",
      "\n",
      "Matched input and pred segments\n",
      "go go we should keep moving\n",
      "vamos, vamos.\n",
      "\n",
      ">go\n",
      "<vamos, vamos.\n",
      "\n",
      "1897 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_16/spa-eng/segments_eng/heroes_s3_16_eng_aligned_eng0093.csv\n",
      "['¡sí!']\n",
      "['sí']\n",
      "syllable_similarity_penalty 2.982456140350877\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['yeah', 'in here']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ sí !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡sí!']\n",
      "\n",
      "Matched input and pred segments\n",
      "yeah in here\n",
      "¡sí!\n",
      "\n",
      ">yeah\n",
      "<¡sí!\n",
      "\n",
      "1957 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_5/spa-eng/segments_eng/heroes_s2_5_eng_aligned_eng0295.csv\n",
      "['¿de', 'acuerdo?']\n",
      "['de', 'acuerdo']\n",
      "syllable_similarity_penalty 3.0701754385964914\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['i', 'sleep in a pub remember', 'you paint']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁¿ de ▁acuerdo ?']\n",
      "\n",
      "Segmented output tokens\n",
      "['¿de acuerdo?']\n",
      "\n",
      "Matched input and pred segments\n",
      "i sleep in a pub remember you paint\n",
      "¿de acuerdo?\n",
      "\n",
      ">i\n",
      "<¿de acuerdo?\n",
      "\n",
      "2479 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_12/spa-eng/segments_eng/heroes_s3_12_eng_aligned_eng0110.csv\n",
      "['¿hola?']\n",
      "['hola']\n",
      "syllable_similarity_penalty 3.1704260651629075\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['hello', 'cheerleader', 'hello']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁¿ hola ?']\n",
      "\n",
      "Segmented output tokens\n",
      "['¿hola?']\n",
      "\n",
      "Matched input and pred segments\n",
      "hello cheerleader hello\n",
      "¿hola?\n",
      "\n",
      ">hello\n",
      "<¿hola?\n",
      "\n",
      "2738 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_4/spa-eng/segments_eng/heroes_s2_4_eng_aligned_eng0379.csv\n",
      "['¡dios!']\n",
      "['dios']\n",
      "syllable_similarity_penalty 3.2017543859649127\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['homicidio', 'right there']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ dios !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡dios!']\n",
      "\n",
      "Matched input and pred segments\n",
      "homicidio right there\n",
      "¡dios!\n",
      "\n",
      ">homicidio\n",
      "<¡dios!\n",
      "\n",
      "2917 - /Users/alp/Movies/heroes/corpus_post/heroes_s3_12/spa-eng/segments_eng/heroes_s3_12_eng_aligned_eng0082.csv\n",
      "['vale.']\n",
      "['vale']\n",
      "syllable_similarity_penalty 2.5438596491228074\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['fine', 'i give']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁vale .']\n",
      "\n",
      "Segmented output tokens\n",
      "['vale.']\n",
      "\n",
      "Matched input and pred segments\n",
      "fine i give\n",
      "vale.\n",
      "\n",
      ">fine\n",
      "<vale.\n",
      "\n",
      "3304 - /Users/alp/Movies/heroes/corpus_post/heroes_s2_5/spa-eng/segments_eng/heroes_s2_5_eng_aligned_eng0191.csv\n",
      "['¡vamos!']\n",
      "['vamos']\n",
      "syllable_similarity_penalty 2.982456140350877\n",
      "Couldn't match all input phrases\n",
      "=======\n",
      "Input segments (from proscript tokens)\n",
      "['come', 'on unlock these cuffs']\n",
      "\n",
      "Segmented output pieces\n",
      "['▁ ¡ vamos !']\n",
      "\n",
      "Segmented output tokens\n",
      "['¡vamos!']\n",
      "\n",
      "Matched input and pred segments\n",
      "come on unlock these cuffs\n",
      "¡vamos!\n",
      "\n",
      ">come\n",
      "<¡vamos!\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#STUDY PARTICULAR DATA SAMPLE\n",
    "#WORKING_SENTENCE_NO = 134\n",
    "for WORKING_SENTENCE_NO in [111, 134, 367, 643, 683, 980, 1242, 1310, 1420, 1732, 1897, 1957, 2479, 2738, 2917, 3304]:\n",
    "    input_pieces = all_input_pieces[WORKING_SENTENCE_NO]\n",
    "    pred_pieces = all_pred_pieces[WORKING_SENTENCE_NO]\n",
    "    input_proscript = all_proscript[WORKING_SENTENCE_NO]\n",
    "    sent_attention_binary = all_attention_binary[WORKING_SENTENCE_NO]\n",
    "    sent_attention_weights = all_attention_weights[WORKING_SENTENCE_NO]\n",
    "    proscript_path = all_proscript_path[WORKING_SENTENCE_NO]\n",
    "\n",
    "    print(\"%i - %s\"%(WORKING_SENTENCE_NO, proscript_path))\n",
    "\n",
    "    # 1. read from processed batch\n",
    "    # input_tokens = all_input_tokens[WORKING_SENTENCE_NO]\n",
    "    # input_token_labels = all_input_token_labels[WORKING_SENTENCE_NO]\n",
    "    # input_piece_labels = all_input_piece_labels[WORKING_SENTENCE_NO]\n",
    "    # input_phrase_structure = all_input_phrase_structure[WORKING_SENTENCE_NO]\n",
    "\n",
    "    # pred_tokens = all_pred_tokens[WORKING_SENTENCE_NO]\n",
    "    # pred_piece_labels = all_pred_piece_labels[WORKING_SENTENCE_NO]\n",
    "    # pred_token_labels = all_pred_token_labels[WORKING_SENTENCE_NO]\n",
    "    # pred_phrase_structure = all_pred_phrase_structure[WORKING_SENTENCE_NO]\n",
    "\n",
    "    # 2. process your own\n",
    "    input_tokens, input_token_labels, input_phrase_structure = get_input_structure(input_proscript, min_segment_pause_interval=MIN_SEGMENT_PAUSE_INTERVAL)\n",
    "    input_piece_labels = token_to_piece_labels(input_pieces, input_tokens, input_token_labels)\n",
    "    #pred_piece_labels = map_piece_labels_naive(input_piece_labels, pred_pieces, sent_attention_binary)\n",
    "    pred_piece_labels, possible_pred_label_seqs, possible_seq_attn_scores = get_pred_labels(input_pieces, input_piece_labels, pred_pieces, sent_attention_weights)\n",
    "    pred_tokens, pred_token_labels = piece_to_token(pred_pieces, pred_piece_labels)\n",
    "    pred_phrase_structure, problem_matching = get_output_phrase_structure(pred_tokens, pred_token_labels, input_phrase_structure)\n",
    "\n",
    "    syllable_similarity_penalty = get_syllable_similarity_penalty(input_phrase_structure, pred_phrase_structure)\n",
    "    print('syllable_similarity_penalty', syllable_similarity_penalty)\n",
    "\n",
    "    if problem_matching:\n",
    "        print(\"Couldn't match all input phrases\")\n",
    "\n",
    "    print(\"=======\")\n",
    "\n",
    "    input_segments_from_tokens = [' '.join(p.tokens) for p in input_phrase_structure]\n",
    "\n",
    "    print(\"Input segments (from proscript tokens)\")\n",
    "    print(input_segments_from_tokens)\n",
    "\n",
    "\n",
    "    pred_segment_pieces = [[] for i in range(max(pred_piece_labels) + 1)]\n",
    "    for t,i in zip(pred_pieces, pred_piece_labels):\n",
    "        pred_segment_pieces[i].append(t)\n",
    "\n",
    "    pred_segments_from_pieces = [' '.join(segment_pieces) for segment_pieces in pred_segment_pieces]\n",
    "\n",
    "    print('\\nSegmented output pieces')\n",
    "    print(pred_segments_from_pieces)\n",
    "\n",
    "    pred_segment_tokens = [[] for i in range(max(pred_token_labels) + 1)]\n",
    "    for t,i in zip(pred_tokens, pred_token_labels):\n",
    "        pred_segment_tokens[i].append(t)\n",
    "\n",
    "    pred_segments_from_tokens = [' '.join(segment_tokens) for segment_tokens in pred_segment_tokens]\n",
    "\n",
    "    print(\"\\nSegmented output tokens\")\n",
    "    print(pred_segments_from_tokens)\n",
    "\n",
    "    #SHOW matched input and prediction segments\n",
    "    print(\"\\nMatched input and pred segments\")\n",
    "    full_input = ' '.join(input_segments_from_tokens)\n",
    "    full_translation = ' '.join(pred_segments_from_tokens)\n",
    "    print(full_input)\n",
    "    print(full_translation)\n",
    "    print()\n",
    "\n",
    "    for input_seg, pred_seg in zip(input_segments_from_tokens, pred_segments_from_tokens):\n",
    "        print(\">\" + input_seg)\n",
    "        print(\"<\" + pred_seg)\n",
    "        print()\n",
    "\n",
    "    #output files for synthesis\n",
    "    OUTPUT_DIR = 'samples'\n",
    "    #1. plain translation string\n",
    "    sample_id = ''.join(os.path.splitext(os.path.basename(proscript_path))[0].split(\"_eng_aligned\")).replace('eng', '').replace('heroes_', '')\n",
    "    translation_output_file = os.path.join(OUTPUT_DIR, sample_id + '.translation.transcript.txt')\n",
    "    full_translation = ' '.join(pred_tokens)\n",
    "    with open (translation_output_file, 'w') as f_txt:\n",
    "        f_txt.write(full_translation)\n",
    "\n",
    "    #2. phrase structure\n",
    "    phrase_output_file = os.path.join(OUTPUT_DIR, sample_id + '.translation.structure.txt')\n",
    "    with open(phrase_output_file, 'w') as f:\n",
    "        for st in pred_phrase_structure:\n",
    "            f.write(\"%f\\t%f\\t%f\\t%s\\n\"%(st.start_time, st.end_time, st.pause_after, st.tokens))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1732\n"
     ]
    }
   ],
   "source": [
    "for i, f in enumerate(all_proscript_path):\n",
    "    if 's3_9' in f and '0061' in f:\n",
    "        print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/Users/alp/Movies/heroes/corpus_post/heroes_s3_6/spa-eng/segments_eng/heroes_s3_6_eng_aligned_eng0096.csv'"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_proscript_path[11]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_phrase_structure, problem_matching = get_output_phrase_structure(pred_tokens, pred_token_labels, input_phrase_structure)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7",
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}