transcription_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "This is a set of scripts to prepare a custom Russian dataset for training punctuator/capitalizer\n",
    "based on BERT architecture using NeMo Python scripts.\n",
    "\n",
    "The corpus is combined from OpenSubtitles' Russian subtitles and transcriptions from Russian\n",
    "Gosduma and Moscow City Duma.\n",
    "\n",
    "Originally it was also planned to use the Lenta news corpus,\n",
    "however a detailed analysis showed that recent news texts may\n",
    "affect the punctuator/capitalizer badly (mostly due direct speech inside\n",
    "the news texts). So it's not recommended to use it.\n",
    "The scripts clean the texts and convert them into NeMo training format (separate files\n",
    "of lowercased words without punctuation and file with punctuation/capitalization labels\n",
    "for each word in the corresponding first file).\n",
    "\n",
    "The source corpuses are not that big, however their quality is quite high. Trying to find other\n",
    "good Russian corpuses did not bring good results so far. For example, social media comments\n",
    "are easily available but extremely illiterate. Mass media tend to include irrelevant info\n",
    "into their news texts (like comments from social networks).\n",
    "Large literature corpuses are barely available and usually have a lot of errors and\n",
    "irregularities due to OCR mistakes. And besides that, using originally written texts as\n",
    "sources for punctuator training may provide irrelevant results on spoken sentences, which\n",
    "differ significantly from written Russian.\n",
    "That's why professionally written transcriptions of speeches (including subtitles) seem\n",
    "to be a good source of training data.\n",
    "\n",
    "One important feature of the dataset is that it tries to concatenate as many sentences into\n",
    "one training line as possible. For that, it appends sentences to the current line\n",
    "until the count of tokens in its tokenized version exceeds the model input. Tokenizing\n",
    "each string each time makes the preparation procedure quite time-consuming. But without\n",
    "it it's almost impossible to train punctuator to place dots (that is, to separate sentences).\n",
    "There are a lot of punctuators in the Internet which are mostly incapable of handling multi-sentence\n",
    "texts, because they were trained on datasets with one sentence per line.\n",
    "\n",
    "Another feature of the dataset is that considers additional labels, comparing to the original\n",
    "NeMo BERT procedure:\n",
    "1. Punctuation labels, including complex modifiers like -, —, … etc.\n",
    "2. Capitalization label for abbreviation - T (meaning \"total uppercase\", all letters in the word).\n",
    "\n",
    "The original NeMo punctuator cannot apply some of the custom labels properly. For example,\n",
    "label \"-\" additionally means that there should be no space between this word and the next one.\n",
    "Label \"—\" additionally means that there should be extra space between the word itself and the \"—\" symbol.\n",
    "Label \"T\" means that all letters should be uppercased, which NeMo algorithm of applying\n",
    "labels does not understand.\n",
    "So it is advised to users of punctuator trained on this dataset to do manual post-processing\n",
    "(including manually applying the labels) and use \"return_labels=True\" parameters during inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-11T07:04:44.864649182Z",
     "start_time": "2023-11-11T07:04:39.669857209Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import calendar\n",
    "import csv\n",
    "import datetime\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import subprocess\n",
    "import tarfile\n",
    "import time\n",
    "import traceback\n",
    "import shutil\n",
    "import urllib\n",
    "from bs4 import BeautifulSoup\n",
    "from nemo.collections.common.tokenizers import AutoTokenizer\n",
    "from typing import List, Dict\n",
    "from urllib.request import urlopen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Constants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:29.891703780Z",
     "start_time": "2023-11-10T22:29:29.877156891Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "RESULT_DATASET_OUTPUT_DIR = os.path.join(os.getcwd(), 'dataset')\n",
    "DATASET_HEADER = ['source_type', 'source_entity', 'text', 'labels']\n",
    "MAX_TOKENS_IN_MODEL_INPUT = 512\n",
    "TOTAL_SET_FILE = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'punctuation_dataset.csv')\n",
    "TRAIN_PERCENTS = 80\n",
    "DEV_PERCENTS = 10\n",
    "\n",
    "GOSDUMA_SOURCE_FOLDER = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_gosduma_source')\n",
    "GOSDUMA_START_DATE = datetime.datetime(1994, 1, 1)\n",
    "GOSDUMA_SUBSET_FILE = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_gosduma_dataset.csv')\n",
    "\n",
    "MOSDUMA_SOURCE_FOLDER = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_mosduma_source')\n",
    "PARAGRAPH_OR_HEADING_PATTERN = re.compile(\"^(p|h[1-6])$\")\n",
    "URL_FETCH_TIMEOUT = 5.0\n",
    "MOSDUMA_SUBSET_FILE = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_mosduma_dataset.csv')\n",
    "\n",
    "SUBTITLES_CORPUS_FOLDER = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'Subtitles', 'home', 'tsha', 'Subtitles')\n",
    "TAGGED_SUBTITLES_FOLDER = os.path.join(SUBTITLES_CORPUS_FOLDER, 'tagged_texts')\n",
    "RAW_SUBTITLES_FOLDER = os.path.join(SUBTITLES_CORPUS_FOLDER, 'texts')\n",
    "SUBTITLES_SUBSET_FILE = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_subtitles_dataset.csv')\n",
    "\n",
    "ROUND_BRACKETS_PATTERN = r'\\([^)]*\\)'\n",
    "ANGLE_BRACKETS_PATTERN = r'\\<[^>]*\\>'\n",
    "SQUARE_BRACKETS_PATTERN = r'\\[.*?\\]'\n",
    "WORD_PATTERN = re.compile(\"[^\\W_](?:[^\\W_]|-|—|:|;|!|\\?|\\.|,|…)*\")\n",
    "\n",
    "LENTA_SOURCE_FOLDER = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_lenta_source')\n",
    "LENTA_SOURCE_FILE_NAME = 'lenta-ru-news.csv'\n",
    "LENTA_TEXT_ONLY_FILE = 'lenta_texts.txt'\n",
    "# Too many news in the original corpus, so we cut some old texts out.\n",
    "LENTA_OBSOLETE_NEWS_COUNT_THRESHOLD = 500000\n",
    "LENTA_SUBSET_FILE = os.path.join(RESULT_DATASET_OUTPUT_DIR, 'ru_lenta_dataset.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Prepare tokenizer to count dataset lines' tokens (as the model has a limited input layer, and each training raw should not exceed it)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:35.179578094Z",
     "start_time": "2023-11-10T22:29:32.529533729Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer(\n",
    "    pretrained_model_name=\"DeepPavlov/rubert-base-cased-conversational\",\n",
    "    vocab_file=None,\n",
    "    use_fast=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method for splitting text into raw uncased string without punctuation and a list of punctuation_capitalization labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:37.407581182Z",
     "start_time": "2023-11-10T22:29:37.350039644Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def make_dataset_entity(text: str) -> Dict[str, str]:\n",
    "    result: Dict[str, str] = dict()\n",
    "    words = text.split()\n",
    "    proper_words_with_punctuation = []\n",
    "    for word in words:\n",
    "        if WORD_PATTERN.match(word):\n",
    "            proper_words_with_punctuation.append(word)\n",
    "    raw_words = []\n",
    "    labels_array = []\n",
    "    for word in proper_words_with_punctuation:\n",
    "        hyphen_words = word.split('-')\n",
    "        hyphen_part_index = 0\n",
    "        for dash_word in hyphen_words:\n",
    "            cleaned_word = dash_word[:]\n",
    "            cleaned_word = re.sub('[.,\\-—:;!?…]+', '', cleaned_word)\n",
    "            if not cleaned_word:\n",
    "                hyphen_part_index += 1\n",
    "                continue\n",
    "            cleaned_word = cleaned_word.lower()\n",
    "\n",
    "            capit_mark = 'U'\n",
    "            if dash_word.upper() == dash_word:\n",
    "                capit_mark = 'T'\n",
    "            elif dash_word.lower() == dash_word:\n",
    "                capit_mark = 'O'\n",
    "\n",
    "            punct_mark = 'O'\n",
    "            if hyphen_part_index < len(hyphen_words) - 1:\n",
    "                punct_mark = '-'\n",
    "            elif dash_word.endswith('...'):\n",
    "                punct_mark = '…'\n",
    "            elif dash_word.endswith('…'):\n",
    "                 punct_mark = '…'\n",
    "            elif dash_word.endswith('.'):\n",
    "                punct_mark = '.'\n",
    "            elif dash_word.endswith(','):\n",
    "                punct_mark = ','\n",
    "            elif dash_word.endswith(':'):\n",
    "                punct_mark = ':'\n",
    "            elif dash_word.endswith(';'):\n",
    "                punct_mark = ';'\n",
    "            elif dash_word.endswith('?!'):\n",
    "                punct_mark = '⁈'\n",
    "            elif dash_word.endswith('!?'):\n",
    "                punct_mark = '⁈'\n",
    "            elif dash_word.endswith('—'):\n",
    "                punct_mark = '—'\n",
    "            elif dash_word.endswith('?'):\n",
    "                punct_mark = '?'\n",
    "            elif dash_word.endswith('!'):\n",
    "                punct_mark = '!'\n",
    "\n",
    "            raw_words.append(cleaned_word)\n",
    "            labels_array.append(f\"{punct_mark}{capit_mark}\")\n",
    "\n",
    "            hyphen_part_index += 1\n",
    "\n",
    "    result['text'] = ' '.join(raw_words)\n",
    "    result['labels'] = ' '.join(labels_array)\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Taiga Subtitles corpus link (around 1 GB): http://bit.ly/2JLeujk , website: https://tatianashavrina.github.io/taiga_site/downloads.html .\n",
    "Place the tar.gx archive into the RESULT_DATASET_OUTPUT_DIR (datasets/ru_punctuator) folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Put the Subtitles corpus from Taiga dataset folder to the folder manually (see the comment above)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Extract the Taiga Subtitles archive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "subtitles_archive = tarfile.open(os.path.join(RESULT_DATASET_OUTPUT_DIR, 'Subtitles.tar.gz'))\n",
    "subtitles_archive.extractall(os.path.join(RESULT_DATASET_OUTPUT_DIR, 'Subtitles'))\n",
    "subtitles_archive.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Delete redundant folders and files from the Subtitles corpus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if os.path.exists(TAGGED_SUBTITLES_FOLDER):\n",
    "    shutil.rmtree(TAGGED_SUBTITLES_FOLDER)\n",
    "for _, film_dirs, _ in os.walk(RAW_SUBTITLES_FOLDER):\n",
    "    for film_folder in film_dirs:\n",
    "        for _, _, subtitle_raw_files in os.walk(os.path.join(RAW_SUBTITLES_FOLDER, film_folder)):\n",
    "            for subtitle_file in subtitle_raw_files:\n",
    "                subtitle_file_path = os.path.join(RAW_SUBTITLES_FOLDER, film_folder, subtitle_file)\n",
    "                if not subtitle_file.endswith(\"ru.txt\"):\n",
    "                    os.remove(subtitle_file_path)\n",
    "                elif os.stat(subtitle_file_path).st_size == 0:\n",
    "                    os.remove(subtitle_file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method for pre-processing a subtitle line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_subtitle_line(line: str) -> str:\n",
    "    processed_line = ''\n",
    "    try:\n",
    "        # The format is like 1\\t00:00:01,200\\t00:00:02,630\\tДоброе утро.\n",
    "        # So we need to skip the tab character 3 times before the actual text begins.\n",
    "        processed_line = line[line.index('\\t') + 1:]\n",
    "        processed_line = processed_line[processed_line.index('\\t') + 1:]\n",
    "        processed_line = processed_line[processed_line.index('\\t') + 1:]\n",
    "        if processed_line.endswith('\\n'):\n",
    "            processed_line = processed_line[:-1]\n",
    "        processed_line = re.sub(ROUND_BRACKETS_PATTERN, '', processed_line)\n",
    "        processed_line = re.sub(SQUARE_BRACKETS_PATTERN, '', processed_line)\n",
    "        if not processed_line:\n",
    "            return ''\n",
    "\n",
    "        if not processed_line or any([processed_line.find(c) != -1 for c in ['<', '>', '♪', '(', ')', '[', ']', '/']]):\n",
    "            print(f\"Skipped line: {processed_line}\")\n",
    "            return ''\n",
    "\n",
    "        processed_line = processed_line.replace('\\\"', '')\n",
    "        processed_line = processed_line.replace('\\'', '')\n",
    "        processed_line = processed_line.replace('«', '')\n",
    "        processed_line = processed_line.replace('»', '')\n",
    "        processed_line = processed_line.replace('%', ' процентов ')\n",
    "        processed_line = processed_line.replace('№', ' номер ')\n",
    "        processed_line = re.sub(' {2,}', ' ', processed_line)\n",
    "        processed_line = processed_line.replace(' ...', '...')\n",
    "        processed_line = processed_line.replace('...', '…')\n",
    "        processed_line = re.sub('\\.{2}', '.', processed_line)\n",
    "        if processed_line.startswith('- '):\n",
    "            processed_line = processed_line[2:]\n",
    "        if processed_line.startswith('— '):\n",
    "            processed_line = processed_line[2:]\n",
    "        if processed_line.startswith('-'):\n",
    "            processed_line = processed_line[1:]\n",
    "        if processed_line.startswith('—'):\n",
    "            processed_line = processed_line[1:]\n",
    "        processed_line = processed_line.replace('. -', '. ')\n",
    "        processed_line = processed_line.replace('? -', '? ')\n",
    "        processed_line = processed_line.replace('! -', '! ')\n",
    "        processed_line = processed_line.replace('. —', '. ')\n",
    "        processed_line = processed_line.replace('? —', '? ')\n",
    "        processed_line = processed_line.replace('! —', '! ')\n",
    "        processed_line = processed_line.replace('!...', '!')\n",
    "        processed_line = processed_line.replace('?...', '?')\n",
    "        processed_line = processed_line.replace('...!', '!')\n",
    "        processed_line = processed_line.replace('...?', '?')\n",
    "        processed_line = processed_line.replace(' - ', '— ')\n",
    "        processed_line = processed_line.replace(' —', '—')\n",
    "        processed_line = re.sub(' {2,}', ' ', processed_line)\n",
    "        return processed_line\n",
    "    except Exception as e:\n",
    "        print(f\"Could not parse a subtitle line {processed_line} due to {e}, {traceback.format_exc()}\")\n",
    "    return ''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method to save a subset of dataset to a CSV file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:45.763310551Z",
     "start_time": "2023-11-10T22:29:45.738326362Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def save_entities_to_csv(dataset_entities: List[Dict[str, str]], file_path: str):\n",
    "    with open(file_path, \"w\") as target_file:\n",
    "        writer = csv.writer(target_file)\n",
    "        writer.writerow(DATASET_HEADER)\n",
    "        for dataset_entity in dataset_entities:\n",
    "            writer.writerow([\n",
    "                dataset_entity['source_type'],\n",
    "                dataset_entity['source_entity'],\n",
    "                dataset_entity['text'],\n",
    "                dataset_entity['labels']\n",
    "            ])\n",
    "    print(f'File {file_path} saved.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Prepare a punctuation/capitalization dataset from Subtitles corpus\n",
    "(split into uncased words and punctuation/capitalization marks)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "subtitles_dataset_entities: List[Dict[str, str]] = []\n",
    "folder_counter = 0\n",
    "for _, film_dirs, _ in os.walk(RAW_SUBTITLES_FOLDER):\n",
    "    for film_folder in film_dirs:\n",
    "        print(f\"Entering {folder_counter} folder {film_folder}\")\n",
    "\n",
    "        for _, _, subtitle_raw_files in os.walk(\n",
    "                os.path.join(RAW_SUBTITLES_FOLDER, film_folder)):\n",
    "            for subtitle_file in subtitle_raw_files:\n",
    "                subtitle_file_path = os.path.join(RAW_SUBTITLES_FOLDER, film_folder,\n",
    "                                                  subtitle_file)\n",
    "                with open(subtitle_file_path) as subtitles_file:\n",
    "                    print(f\"Entering file {subtitle_file_path}\")\n",
    "                    lines_candidates = []\n",
    "                    last_lines_not_ended_with_stop = []\n",
    "                    for subtitle_line in subtitles_file:\n",
    "                        preprocessed_line = preprocess_subtitle_line(subtitle_line)\n",
    "                        if not preprocessed_line:\n",
    "                            continue\n",
    "                        entity_text_candidate = ' '.join(lines_candidates) + \\\n",
    "                                                ' ' + preprocessed_line\n",
    "                        entity_tokens = tokenizer.text_to_tokens(entity_text_candidate)\n",
    "                        if len(entity_tokens) < MAX_TOKENS_IN_MODEL_INPUT:\n",
    "                            if preprocessed_line.endswith('.') \\\n",
    "                                    or preprocessed_line.endswith('…') \\\n",
    "                                    or preprocessed_line.endswith('!') \\\n",
    "                                    or preprocessed_line.endswith('?'):\n",
    "                                for line_not_ended_with_stop in last_lines_not_ended_with_stop:\n",
    "                                    lines_candidates.append(line_not_ended_with_stop)\n",
    "                                last_lines_not_ended_with_stop = []\n",
    "                                lines_candidates.append(preprocessed_line)\n",
    "                            else:\n",
    "                                last_lines_not_ended_with_stop.append(preprocessed_line)\n",
    "                        else:\n",
    "                            text_for_entity = ' '.join(lines_candidates)\n",
    "                            lines_candidates = last_lines_not_ended_with_stop[:]\n",
    "                            lines_candidates.append(preprocessed_line)\n",
    "                            last_lines_not_ended_with_stop = []\n",
    "                            if text_for_entity:\n",
    "                                dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                                if dataset_entity:\n",
    "                                    dataset_entity['source_type'] = 'subtitles'\n",
    "                                    dataset_entity['source_entity'] = f'{film_folder}; {subtitle_file}.'\n",
    "                                    subtitles_dataset_entities.append(dataset_entity)\n",
    "\n",
    "                    # If anything left in the string buffer and the token limit is not\n",
    "                    # reached yet - add it to the dataset forcefully.\n",
    "                    if lines_candidates:\n",
    "                        text_for_entity = ' '.join(lines_candidates)\n",
    "                        if text_for_entity:\n",
    "                            dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                            if dataset_entity:\n",
    "                                dataset_entity['source_type'] = 'subtitles'\n",
    "                                dataset_entity['source_entity'] = f'{film_folder}; {subtitle_file}.'\n",
    "                                subtitles_dataset_entities.append(dataset_entity)\n",
    "\n",
    "        folder_counter += 1\n",
    "\n",
    "save_entities_to_csv(subtitles_dataset_entities, SUBTITLES_SUBSET_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method to detect the end of speaker’s introduction in the transcription line, considering that the introduction has at least one comma\n",
    "(for complex introductions like “Иванов Иван Иванович, председатель законодательного собрания.”)\n",
    "Since the introduction can be in the same line as the actual speech, consider the end of it at the nearest dot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:50.989909976Z",
     "start_time": "2023-11-10T22:29:50.986151211Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def detect_introduction_with_comma(line: str) -> int:\n",
    "    line_after_comma = line[line.index(',') + 1:]\n",
    "    sentence_separator_position = line_after_comma.find('. ')\n",
    "    if sentence_separator_position != -1:\n",
    "        return line.index(',') + 1 + sentence_separator_position + len('. ')\n",
    "    else:\n",
    "        return len(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method to find the end of an introduction of the speaker in the transcription. A typical introduction is\n",
    "<last name> <first initial>.(<last initial>.)(, <title>.)\n",
    "And a hardcoded introduction is \"Председательствующий\", which means the chairman.\n",
    "If the beginning of the string matches the introduction pattern, then returns the position of its last symbol.\n",
    "When preparing punctuation dataset, the introduction in the lines should not be considered, as it's not a part of the actual speech."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:51.906630313Z",
     "start_time": "2023-11-10T22:29:51.860824740Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def detect_introduction(line: str) -> int:\n",
    "    if line.find(' ') == -1:\n",
    "        return 0\n",
    "\n",
    "    if line.startswith('Председательствующий. '):\n",
    "        return len('Председательствующий. ')\n",
    "\n",
    "    first_word = line[:line.index(' ')]\n",
    "    if not first_word or first_word.capitalize() != first_word \\\n",
    "            or not re.match(r'[^\\W\\d_](?:[^\\W\\d_]| |-)*[^\\W\\d_]', first_word, re.UNICODE):\n",
    "        return 0\n",
    "\n",
    "    line_after_first_word = line[line.index(' ') + 1:]\n",
    "    if not line_after_first_word or line_after_first_word.find(' ') == -1:\n",
    "        return 0\n",
    "    second_word = line_after_first_word[:line_after_first_word.index(' ')]\n",
    "    if not second_word:\n",
    "        return 0\n",
    "    if second_word[0].upper() == second_word[0] \\\n",
    "            and re.match(r'[^\\W\\d_]', second_word[0], re.UNICODE):\n",
    "        if len(second_word) > 1 and second_word[1] == '.':\n",
    "            if len(second_word) > 3:\n",
    "                if second_word[2].upper() == second_word[2] \\\n",
    "                        and re.match(r'[^\\W\\d_]', second_word[2], re.UNICODE) \\\n",
    "                        and second_word[3] == '.':\n",
    "                    if len(second_word) == 5:\n",
    "                        if second_word[4] == ',':\n",
    "                            return detect_introduction_with_comma(line)\n",
    "                    elif len(second_word) == 4:\n",
    "                        return len(first_word) + 1 + 4\n",
    "            elif len(second_word) == 2:\n",
    "                line_after_second_word = line_after_first_word[\n",
    "                                         line_after_first_word.index(' ') + 1:]\n",
    "                if not line_after_second_word or line_after_second_word.find(' ') == -1:\n",
    "                    return 0\n",
    "                third_word = line_after_second_word[:line_after_second_word.index(' ')]\n",
    "\n",
    "                if not third_word:\n",
    "                    return 0\n",
    "                if len(third_word) == 2 \\\n",
    "                        and third_word[0].upper() == third_word[0] \\\n",
    "                        and re.match(r'[^\\W\\d_]', third_word[0], re.UNICODE) \\\n",
    "                        and third_word[1] == '.':\n",
    "                    return len(first_word) + 1 + len(second_word) + 1 + len(third_word) + 1\n",
    "                elif len(third_word) == 3 \\\n",
    "                        and third_word[0].upper() == third_word[0] \\\n",
    "                        and re.match(r'[^\\W\\d_]', third_word[0], re.UNICODE) \\\n",
    "                        and third_word[1] == '.' \\\n",
    "                        and third_word[2] == ',':\n",
    "                    return detect_introduction_with_comma(line)\n",
    "                else:\n",
    "                    return len(first_word) + 1 + len(second_word) + 1\n",
    "            else:\n",
    "                if second_word[3] == ',':\n",
    "                    return detect_introduction_with_comma(line)\n",
    "\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Preprocess a line of transcription or regular written text, removes punctuation signs that cannot be handled by punctuator.\n",
    "If the line contains indicators that it is malformed (uneven count of brackets etc.), then returns nothing not to consider\n",
    "it for the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:53.179586198Z",
     "start_time": "2023-11-10T22:29:53.122874440Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def clean_line(line: str) -> str:\n",
    "    if not line:\n",
    "        return ''\n",
    "\n",
    "    if not line or any([line.find(c) != -1 for c in ['<', '>', '♪', '(', ')', '[', ']']]):\n",
    "        print(f\"Skipped line: {line}\")\n",
    "        return ''\n",
    "\n",
    "    line = line.replace('/', ' ')\n",
    "    line = line.replace('%', ' процентов ')\n",
    "    line = line.replace('№', ' номер ')\n",
    "    line = re.sub(r'(\\d),(\\d)', r'\\1 и \\2', line)\n",
    "    line = line.replace('\\\"', '')\n",
    "    line = line.replace('\\'', '')\n",
    "    line = line.replace('«', '')\n",
    "    line = line.replace('»', '')\n",
    "    line = re.sub(r' {2,}', ' ', line)\n",
    "    line = line.replace(' ...', '...')\n",
    "    line = line.replace('...', '…')\n",
    "    line = re.sub(r'\\.{2}', '.', line)\n",
    "\n",
    "    if line.startswith('- '):\n",
    "        line = line[2:]\n",
    "    if line.startswith('— '):\n",
    "        line = line[2:]\n",
    "    if line.startswith('-'):\n",
    "        line = line[1:]\n",
    "    if line.startswith('—'):\n",
    "        line = line[1:]\n",
    "    line = line.replace('. -', '. ')\n",
    "    line = line.replace('? -', '? ')\n",
    "    line = line.replace('! -', '! ')\n",
    "    line = line.replace('. —', '. ')\n",
    "    line = line.replace('? —', '? ')\n",
    "    line = line.replace('! —', '! ')\n",
    "    line = line.replace('!...', '!')\n",
    "    line = line.replace('?...', '?')\n",
    "    line = line.replace('...!', '!')\n",
    "    line = line.replace('...?', '?')\n",
    "    line = line.replace(' - ', '— ')\n",
    "    line = line.replace(' —', '—')\n",
    "    line = re.sub(' {2,}', ' ', line)\n",
    "    line = line.strip()\n",
    "\n",
    "    if not line.endswith('.') \\\n",
    "            and not line.endswith('!') \\\n",
    "            and not line.endswith('?'):\n",
    "        line = line + '.'\n",
    "\n",
    "    return line"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Method to split a line into sentences seperated by dots. Each sentence if then cleared from unwanted symbols,\n",
    "and can be completelly omitted if it looks malformed.\n",
    "The split is needed because a line of transcription can be way longer than the model input limit, so we cannot\n",
    "train the punncuator on such long lines.\n",
    "Please note that we don't use this method for the Subtitles corpus, since it has only very short line,\n",
    "and also because the preprocessing looks a bit different for it due to another syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:54.424646701Z",
     "start_time": "2023-11-10T22:29:54.373883329Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def get_cleaned_lines(line: str) -> List[str]:\n",
    "    line = re.sub(ROUND_BRACKETS_PATTERN, '', line)\n",
    "    line = re.sub(SQUARE_BRACKETS_PATTERN, '', line)\n",
    "    line = re.sub(ANGLE_BRACKETS_PATTERN, '', line)\n",
    "\n",
    "    processed_lines = line.split('. ')\n",
    "\n",
    "    result_lines = []\n",
    "    for processed_line in processed_lines:\n",
    "        processed_line = clean_line(processed_line)\n",
    "        if not processed_line:\n",
    "            continue\n",
    "\n",
    "        result_lines.append(processed_line)\n",
    "    return result_lines"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Preprocesses a line of Gosduma's transcription. Splits it by sentences, removes unnecessary symbols,\n",
    "omits malformed sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:29:55.769220888Z",
     "start_time": "2023-11-10T22:29:55.763510678Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_gosduma_transcription_line(line: str) -> List[str]:\n",
    "    if not line or line == '\\n':\n",
    "        return []\n",
    "    if line.endswith('\\n'):\n",
    "        line = line[:-1]\n",
    "    try:\n",
    "        index_after_introduction = detect_introduction(line)\n",
    "        if index_after_introduction >= len(line):\n",
    "            return []\n",
    "\n",
    "        line_after_introduction = line[index_after_introduction:]\n",
    "\n",
    "        return get_cleaned_lines(line_after_introduction)\n",
    "    except Exception as e:\n",
    "        print(f\"Could not parse a subtitle line {line} due to {e}, {traceback.format_exc()}\")\n",
    "    return []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Prepare raw data for the Gosduma sub-dataset. Download all available transcriptions from the\n",
    "official website and place it to respective local folders by month-year pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "start_date = GOSDUMA_START_DATE\n",
    "if not os.path.exists(GOSDUMA_SOURCE_FOLDER):\n",
    "    os.makedirs(GOSDUMA_SOURCE_FOLDER)\n",
    "next_year = int(datetime.date.today().year) + 1\n",
    "while start_date.year < next_year:\n",
    "    month_subfolder = os.path.join(GOSDUMA_SOURCE_FOLDER, f\"{start_date.month}_{start_date.year}\")\n",
    "    if not os.path.exists(month_subfolder):\n",
    "        os.makedirs(month_subfolder)\n",
    "\n",
    "    start_date = start_date.replace(day=1)\n",
    "    last_day_in_month = calendar.monthrange(start_date.year, start_date.month)[1]\n",
    "    end_date = start_date.replace(day=last_day_in_month)\n",
    "\n",
    "    transcription_url_prefix = \"http://transcript.duma.gov.ru\"\n",
    "    transcriptions_list_url = f\"http://transcript.duma.gov.ru/search/?sessid=0&doctype=3&dt_start={start_date.strftime('%d.%m.%Y')}&dt_end={end_date.strftime('%d.%m.%Y')}&phrase1=\"\n",
    "    print(f\"Transcription list URL: {transcriptions_list_url}\")\n",
    "    list_page = urlopen(transcriptions_list_url)\n",
    "    list_soup = BeautifulSoup(list_page, 'html.parser')\n",
    "    transcription_hrefs = []\n",
    "    for a in list_soup.find_all('a', href=True):\n",
    "        if a['href'].startswith('/node/'):\n",
    "            transcription_hrefs.append(a['href'])\n",
    "\n",
    "    for transcription_href in transcription_hrefs:\n",
    "        full_transcription_url = transcription_url_prefix + transcription_href\n",
    "        transcription_page = urlopen(full_transcription_url)\n",
    "        transcription_soup = BeautifulSoup(transcription_page, 'html.parser')\n",
    "\n",
    "        header_div = transcription_soup.find_all(\"div\", {\"class\": \"header-bord\"})[0]\n",
    "        header_text = header_div.get_text().split('\\n')[1]\n",
    "\n",
    "        content_div = transcription_soup.find_all(\"div\", {\"id\": \"selectable-content\"})[0]\n",
    "        content_text = content_div.get_text()\n",
    "\n",
    "        with open(os.path.join(month_subfolder, f'{header_text}.txt'), 'w') as source_text_file:\n",
    "            source_text_file.write(content_text)\n",
    "            source_text_file.close()\n",
    "        # Keep web-crawling responsible by limiting the requests count per minute.\n",
    "        time.sleep(1)\n",
    "\n",
    "    if start_date.month == 12:\n",
    "        start_date = start_date.replace(year=start_date.year + 1)\n",
    "        start_date = start_date.replace(month=1)\n",
    "    else:\n",
    "        start_date = start_date.replace(month=start_date.month + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Splits the Gosduma raw dataset into uncased words and punctuation/capitalization labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "gosduma_dataset_entities: List[Dict[str, str]] = []\n",
    "folder_counter = 0\n",
    "for _, month_dirs, _ in os.walk(GOSDUMA_SOURCE_FOLDER):\n",
    "    for month_dir in month_dirs:\n",
    "        print(f\"Entering {folder_counter} folder {month_dir}. Count of entities so far: {len(gosduma_dataset_entities)}\")\n",
    "\n",
    "        for _, _, month_raw_files in os.walk(\n",
    "                os.path.join(GOSDUMA_SOURCE_FOLDER, month_dir)):\n",
    "            for transcription_raw_file in month_raw_files:\n",
    "                transcription_file_path = os.path.join(GOSDUMA_SOURCE_FOLDER, month_dir,\n",
    "                                                       transcription_raw_file)\n",
    "                with open(transcription_file_path) as transcription_file:\n",
    "                    print(f\"Entering file {transcription_file_path}\")\n",
    "                    lines_candidates = []\n",
    "                    last_lines_not_ended_with_stop = []\n",
    "                    for transcription_line in transcription_file:\n",
    "                        preprocessed_lines = preprocess_gosduma_transcription_line(transcription_line)\n",
    "                        if not preprocessed_lines:\n",
    "                            continue\n",
    "\n",
    "                        for preprocessed_line in preprocessed_lines:\n",
    "                            if not preprocessed_line:\n",
    "                                continue\n",
    "\n",
    "                            entity_text_candidate = ' '.join(lines_candidates) + \\\n",
    "                                                    ' ' + preprocessed_line\n",
    "                            entity_tokens = tokenizer.text_to_tokens(entity_text_candidate)\n",
    "                            if len(entity_tokens) < MAX_TOKENS_IN_MODEL_INPUT:\n",
    "                                if preprocessed_line.endswith('.') \\\n",
    "                                        or preprocessed_line.endswith('…') \\\n",
    "                                        or preprocessed_line.endswith('!') \\\n",
    "                                        or preprocessed_line.endswith('?'):\n",
    "                                    for line_not_ended_with_stop in last_lines_not_ended_with_stop:\n",
    "                                        lines_candidates.append(line_not_ended_with_stop)\n",
    "                                    last_lines_not_ended_with_stop = []\n",
    "                                    lines_candidates.append(preprocessed_line)\n",
    "                                else:\n",
    "                                    last_lines_not_ended_with_stop.append(preprocessed_line)\n",
    "                            else:\n",
    "                                text_for_entity = ' '.join(lines_candidates)\n",
    "                                lines_candidates = last_lines_not_ended_with_stop[:]\n",
    "                                lines_candidates.append(preprocessed_line)\n",
    "                                last_lines_not_ended_with_stop = []\n",
    "                                if text_for_entity:\n",
    "                                    dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                                    if dataset_entity:\n",
    "                                        dataset_entity['source_type'] = 'gosduma'\n",
    "                                        dataset_entity['source_entity'] = f'{transcription_raw_file}.'\n",
    "                                        gosduma_dataset_entities.append(dataset_entity)\n",
    "\n",
    "                    # If anything left in the string buffer and the token limit is not\n",
    "                    # reached yet - add it to the dataset forcefully.\n",
    "                    if lines_candidates:\n",
    "                        text_for_entity = ' '.join(lines_candidates)\n",
    "                        if text_for_entity:\n",
    "                            dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                            if dataset_entity:\n",
    "                                dataset_entity['source_type'] = 'gosduma'\n",
    "                                dataset_entity['source_entity'] = f'{transcription_raw_file}.'\n",
    "                                gosduma_dataset_entities.append(dataset_entity)\n",
    "\n",
    "        folder_counter += 1\n",
    "\n",
    "save_entities_to_csv(gosduma_dataset_entities, GOSDUMA_SUBSET_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:27:58.595970381Z",
     "start_time": "2023-11-10T22:27:40.342463140Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Preprocesses regular text line of speech (without speaker's introduction). Suitable for Mosduma and Lenta datasets.\n",
    "Splits the lines by sentences, deleted unnecessary symbols, omits malformed sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:30:00.691221954Z",
     "start_time": "2023-11-10T22:30:00.684671005Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_regular_line(line: str) -> List[str]:\n",
    "    if not line or line == '\\n':\n",
    "        return []\n",
    "    if line.endswith('\\n'):\n",
    "        line = line[:-1]\n",
    "    try:\n",
    "        return get_cleaned_lines(line)\n",
    "    except Exception as e:\n",
    "        print(f\"Could not parse a subtitle line {line} due to {e}, {traceback.format_exc()}\")\n",
    "    return []"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Preprocesses a line from Mosduma transcription."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-10T22:30:01.957234063Z",
     "start_time": "2023-11-10T22:30:01.948700022Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_mosduma_transcription_line(line: str) -> List[str]:\n",
    "    return preprocess_regular_line(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Download transcriptions from the official MosGorDuma website, convert doc files into HTML, then parse the raw transcription part\n",
    "from it and save locally.\n",
    "Converting from RTF/DOCX to HTML is needed because speaker's introductions in Mosduma transcriptions can be detected by the\n",
    "bold formatting, which would be lost in case of straightforward doc->txt conversion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-11T07:02:20.509640738Z",
     "start_time": "2023-11-11T07:02:20.238862071Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "if not os.path.exists(MOSDUMA_SOURCE_FOLDER):\n",
    "    os.makedirs(MOSDUMA_SOURCE_FOLDER)\n",
    "page_number = 1\n",
    "last_page_doc_urls: List[str] = []\n",
    "while True:\n",
    "    page_subfolder = os.path.join(MOSDUMA_SOURCE_FOLDER, f\"{page_number}\")\n",
    "    if os.path.exists(page_subfolder):\n",
    "        page_number += 1\n",
    "        continue\n",
    "\n",
    "    transcriptions_list_url = f\"https://duma.mos.ru/ru/60/conference?date%5Bfrom%5D=0&date%5Bto%5D=0&sort=date-desc&page={page_number}#results\"\n",
    "    print(f\"Transcription list URL: {transcriptions_list_url}\")\n",
    "    ssl_browser_req = urllib.request.Request(\n",
    "        transcriptions_list_url,\n",
    "        data=None,\n",
    "        headers={\n",
    "            'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.47 Safari/537.36',\n",
    "            'Sec-Ch-Ua': '\"Chromium\";v=\"115\", \"Not/A)Brand\";v=\"99\"',\n",
    "            'Sec-Ch-Ua-Mobile': '?0',\n",
    "            'Sec-Ch-Ua-Platform': '\"Linux\"',\n",
    "            'Sec-Fetch-Dest': 'document',\n",
    "            'Sec-Fetch-Mode': 'navigate',\n",
    "            'Sec-Fetch-Site': 'none',\n",
    "            'Sec-Fetch-User': '?1',\n",
    "            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7'\n",
    "        }\n",
    "    )\n",
    "    # MosGorDuma website is prone to instabilities, so often retries are needed for most of pages.\n",
    "    while True:\n",
    "        successful_read = False\n",
    "        try:\n",
    "            list_page = urlopen(ssl_browser_req, timeout=URL_FETCH_TIMEOUT).read()\n",
    "            successful_read = True\n",
    "        except Exception:\n",
    "            print(f\"Exception while reading URL: {transcriptions_list_url} ; trace: {traceback.format_exc()}\")\n",
    "        if successful_read:\n",
    "            break\n",
    "\n",
    "    if not os.makedirs(page_subfolder):\n",
    "        os.path.exists(page_subfolder)\n",
    "\n",
    "    list_soup = BeautifulSoup(list_page, 'html.parser')\n",
    "    transcription_hrefs = []\n",
    "    for a in list_soup.find_all('a', href=True):\n",
    "        span = a.find('span')\n",
    "        if span and span.text and span.text.startswith('Стенограмма '):\n",
    "            transcription_hrefs.append(a['href'])\n",
    "\n",
    "    if not transcription_hrefs:\n",
    "        break\n",
    "    # Encountering a doc URL the same as in one of the previous pages means that the last\n",
    "    # page was passed already, and now the website just returns us the content of the last page,\n",
    "    # hence we need to exit the cylce.\n",
    "    # This is a \"feature\" of MosGorDuma's site only.\n",
    "    if any([(transcription_href in last_page_doc_urls) for transcription_href in transcription_hrefs]):\n",
    "        break\n",
    "\n",
    "    for transcription_href in transcription_hrefs:\n",
    "        transcription_request = urllib.request.Request(transcription_href, method='HEAD')\n",
    "        transcription_request_info = urllib.request.urlopen(transcription_request)\n",
    "        file_name = transcription_request_info.info().get_filename()\n",
    "        print(f\"Remote file name: {file_name}\")\n",
    "        raw_transcription_path = os.path.join(page_subfolder, file_name)\n",
    "        # MosGorDuma website is prone to instabilities, so often retries are needed for most of pages.\n",
    "        while True:\n",
    "            successful_download = False\n",
    "            try:\n",
    "                urllib.request.urlretrieve(transcription_href, raw_transcription_path)\n",
    "                successful_download = True\n",
    "            except Exception:\n",
    "                print(f\"Exception while downloading file from URL: {transcription_href} ; trace: {traceback.format_exc()}\")\n",
    "            if successful_download:\n",
    "                break\n",
    "\n",
    "        file_name_wo_extension = file_name[:file_name.rindex('.')]\n",
    "        html_extension = 'html'\n",
    "\n",
    "        subprocess.run(f\"soffice --convert-to {html_extension} {raw_transcription_path} --outdir {page_subfolder}\",\n",
    "                       shell=True)\n",
    "        os.remove(raw_transcription_path)\n",
    "\n",
    "        text_lines = []\n",
    "        html_file_path = os.path.join(page_subfolder, f'{file_name_wo_extension}.{html_extension}')\n",
    "        with open(html_file_path) as html_file:\n",
    "            transcription_soup = BeautifulSoup(html_file, 'html.parser')\n",
    "            start_paragraph = None\n",
    "            for paragraph in transcription_soup.find_all(PARAGRAPH_OR_HEADING_PATTERN):\n",
    "                if paragraph.get_text().replace('\\n', '').replace(' ', '').replace('\\t', '').strip() == \"СТЕНОГРАММА\":\n",
    "                    start_paragraph = paragraph\n",
    "                    print(f\"{file_name_wo_extension}. Start paragraph is СТЕНОГРАММА.\")\n",
    "                    break\n",
    "            if not start_paragraph:\n",
    "                for paragraph in transcription_soup.find_all('p'):\n",
    "                    if paragraph.get_text().replace('\\n', ' ').strip().startswith('файл '):\n",
    "                        start_paragraph = paragraph\n",
    "                        print(f\"{file_name_wo_extension}. Start paragraph is файл: {paragraph.get_text()}.\")\n",
    "                        break\n",
    "            if not start_paragraph:\n",
    "                for paragraph in transcription_soup.find_all(PARAGRAPH_OR_HEADING_PATTERN):\n",
    "                    if paragraph.get_text().replace('\\n', ' ').strip().startswith('Материалы к протоколу заседания') \\\n",
    "                            or paragraph.get_text().replace('\\n', ' ').strip().startswith('Материалы к протоколу внеочередного'):\n",
    "                        materials_paragraph = paragraph\n",
    "                        for paragraph_candidate in materials_paragraph.find_all_next(PARAGRAPH_OR_HEADING_PATTERN):\n",
    "                            if paragraph_candidate.get_text().replace('\\n', ' ').strip().startswith(\n",
    "                                    'Заседание №') \\\n",
    "                                    or paragraph_candidate.get_text().replace('\\n', ' ').strip().startswith(\n",
    "                                    'Внеочередное заседание №') \\\n",
    "                                    or paragraph_candidate.get_text().replace('\\n', ' ').strip().startswith(\n",
    "                                    'Утреннее заседание №') \\\n",
    "                                    or paragraph_candidate.get_text().replace('\\n', ' ').strip().startswith(\n",
    "                                    'Вечернее заседание №'):\n",
    "                                start_paragraph = paragraph_candidate\n",
    "                                print(f\"{file_name_wo_extension}. Start paragraph is Заседание №: {start_paragraph.get_text()}.\")\n",
    "                                break\n",
    "                    if start_paragraph:\n",
    "                        break\n",
    "\n",
    "            if not start_paragraph:\n",
    "                print(f\"page_number {page_number}, file {file_name_wo_extension}, \"\n",
    "                        f\"could not find start of transcription.\")\n",
    "                html_file.close()\n",
    "                os.remove(html_file_path)\n",
    "                continue\n",
    "\n",
    "            for paragraph in start_paragraph.find_all_next('p'):\n",
    "                stripped_text = paragraph.get_text().strip().replace('\\n', ' ')\n",
    "                if stripped_text.startswith('файл '):\n",
    "                    continue\n",
    "                is_all_bold = False\n",
    "                for bold in paragraph.find_all('b'):\n",
    "                    if stripped_text == bold.get_text().strip().replace('\\n', ' '):\n",
    "                        is_all_bold = True\n",
    "                        break\n",
    "\n",
    "                # Adds a line to the result, except introductions, which are detected\n",
    "                # by bold formatting.\n",
    "                if stripped_text and not is_all_bold:\n",
    "                    text_lines.append(stripped_text)\n",
    "\n",
    "            html_file.close()\n",
    "            os.remove(html_file_path)\n",
    "\n",
    "        with open(os.path.join(page_subfolder, f'{file_name_wo_extension}.txt'), 'w') as source_text_file:\n",
    "            source_text_file.write('\\n'.join(text_lines))\n",
    "            source_text_file.close()\n",
    "        # Keep web-crawling responsible by limiting the requests count per minute.\n",
    "        time.sleep(1)\n",
    "\n",
    "    last_page_doc_urls = transcription_hrefs\n",
    "\n",
    "    page_number += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Split the Mosduma raw dataset into uncased words and punctuation/capitalization labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-11T00:15:25.300592859Z",
     "start_time": "2023-11-10T22:32:33.181407753Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "mosduma_dataset_entities: List[Dict[str, str]] = []\n",
    "folder_counter = 0\n",
    "for _, page_dirs, _ in os.walk(MOSDUMA_SOURCE_FOLDER):\n",
    "    for page_dir in page_dirs:\n",
    "        print(f\"Entering {folder_counter} folder {page_dir}\")\n",
    "\n",
    "        for _, _, page_raw_files in os.walk(\n",
    "                os.path.join(MOSDUMA_SOURCE_FOLDER, page_dir)):\n",
    "            for transcription_raw_file in page_raw_files:\n",
    "                transcription_file_path = os.path.join(MOSDUMA_SOURCE_FOLDER, page_dir,\n",
    "                                                       transcription_raw_file)\n",
    "                with open(transcription_file_path) as transcription_file:\n",
    "                    print(f\"Entering file {transcription_file_path}\")\n",
    "                    lines_candidates = []\n",
    "                    last_lines_not_ended_with_stop = []\n",
    "                    for transcription_line in transcription_file:\n",
    "                        preprocessed_lines = preprocess_mosduma_transcription_line(\n",
    "                            transcription_line)\n",
    "                        if not preprocessed_lines:\n",
    "                            continue\n",
    "\n",
    "                        for preprocessed_line in preprocessed_lines:\n",
    "                            if not preprocessed_line:\n",
    "                                continue\n",
    "\n",
    "                            entity_text_candidate = ' '.join(lines_candidates) + \\\n",
    "                                                    ' ' + preprocessed_line\n",
    "                            entity_tokens = tokenizer.text_to_tokens(entity_text_candidate)\n",
    "                            if len(entity_tokens) < MAX_TOKENS_IN_MODEL_INPUT:\n",
    "                                if preprocessed_line.endswith('.') \\\n",
    "                                        or preprocessed_line.endswith('…') \\\n",
    "                                        or preprocessed_line.endswith('!') \\\n",
    "                                        or preprocessed_line.endswith('?'):\n",
    "                                    for line_not_ended_with_stop in last_lines_not_ended_with_stop:\n",
    "                                        lines_candidates.append(line_not_ended_with_stop)\n",
    "                                    last_lines_not_ended_with_stop = []\n",
    "                                    lines_candidates.append(preprocessed_line)\n",
    "                                else:\n",
    "                                    last_lines_not_ended_with_stop.append(preprocessed_line)\n",
    "                            else:\n",
    "                                text_for_entity = ' '.join(lines_candidates)\n",
    "                                lines_candidates = last_lines_not_ended_with_stop[:]\n",
    "                                lines_candidates.append(preprocessed_line)\n",
    "                                last_lines_not_ended_with_stop = []\n",
    "                                if text_for_entity:\n",
    "                                    dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                                    if dataset_entity:\n",
    "                                        dataset_entity['source_type'] = 'mosduma'\n",
    "                                        dataset_entity['source_entity'] = f'{transcription_raw_file}'\n",
    "                                        mosduma_dataset_entities.append(dataset_entity)\n",
    "\n",
    "                    # If anything left in the string buffer and the token limit is not\n",
    "                    # reached yet - add it to the dataset forcefully.\n",
    "                    if lines_candidates:\n",
    "                        text_for_entity = ' '.join(lines_candidates)\n",
    "                        if text_for_entity:\n",
    "                            dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                            if dataset_entity:\n",
    "                                dataset_entity['source_type'] = 'mosduma'\n",
    "                                dataset_entity['source_entity'] = f'{transcription_raw_file}'\n",
    "                                mosduma_dataset_entities.append(dataset_entity)\n",
    "\n",
    "        folder_counter += 1\n",
    "\n",
    "save_entities_to_csv(mosduma_dataset_entities, MOSDUMA_SUBSET_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Preprocesses a line from a news corpus."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def preprocess_news_line(line: str) -> List[str]:\n",
    "    return preprocess_regular_line(line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Makes a raw dataset from the downloaded Lenta corpus.\n",
    "Skips some first lines in order to keep the result dataset more balanced\n",
    "(more spoken sentences, less written).\n",
    "The source archive can be found here: https://github.com/yutkin/Lenta.Ru-News-Dataset/releases\n",
    "Should be downloaded, extracted and placed to the LENTA_SOURCE_FOLDER folder as a CSV file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "with open(os.path.join(LENTA_SOURCE_FOLDER, LENTA_SOURCE_FILE_NAME), 'r') \\\n",
    "        as news_file:\n",
    "    with open(os.path.join(LENTA_SOURCE_FOLDER, LENTA_TEXT_ONLY_FILE),\n",
    "              'w') as lenta_texts_file:\n",
    "        reader = csv.reader(news_file)\n",
    "        news_counter = 0\n",
    "        while news_counter < LENTA_OBSOLETE_NEWS_COUNT_THRESHOLD:\n",
    "            next(reader)\n",
    "            news_counter += 1\n",
    "\n",
    "        while True:\n",
    "            news_line = next(reader, None)\n",
    "\n",
    "            if not news_line:\n",
    "                break\n",
    "\n",
    "            lenta_texts_file.write(news_line[2] + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare CSV with Lenta dataset entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "lenta_dataset_entities: List[Dict[str, str]] = []\n",
    "with open(os.path.join(LENTA_SOURCE_FOLDER, LENTA_TEXT_ONLY_FILE)) as lenta_texts_file:\n",
    "    dataset_entities: List[Dict[str, str]] = []\n",
    "    lines_candidates = []\n",
    "    last_lines_not_ended_with_stop = []\n",
    "    for line in lenta_texts_file:\n",
    "        preprocessed_lines = preprocess_news_line(line)\n",
    "        if not preprocessed_lines:\n",
    "            continue\n",
    "\n",
    "        for preprocessed_line in preprocessed_lines:\n",
    "            if not preprocessed_line:\n",
    "                continue\n",
    "\n",
    "            entity_text_candidate = ' '.join(lines_candidates) + ' ' + preprocessed_line\n",
    "            entity_tokens = tokenizer.text_to_tokens(entity_text_candidate)\n",
    "            if len(entity_tokens) < MAX_TOKENS_IN_MODEL_INPUT:\n",
    "                if preprocessed_line.endswith('.') \\\n",
    "                        or preprocessed_line.endswith('…') \\\n",
    "                        or preprocessed_line.endswith('!') \\\n",
    "                        or preprocessed_line.endswith('?'):\n",
    "                    for line_not_ended_with_stop in last_lines_not_ended_with_stop:\n",
    "                        lines_candidates.append(line_not_ended_with_stop)\n",
    "                    last_lines_not_ended_with_stop = []\n",
    "                    lines_candidates.append(preprocessed_line)\n",
    "                else:\n",
    "                    last_lines_not_ended_with_stop.append(preprocessed_line)\n",
    "            else:\n",
    "                text_for_entity = ' '.join(lines_candidates)\n",
    "                lines_candidates = last_lines_not_ended_with_stop[:]\n",
    "                lines_candidates.append(preprocessed_line)\n",
    "                last_lines_not_ended_with_stop = []\n",
    "                if text_for_entity:\n",
    "                    dataset_entity = make_dataset_entity(text_for_entity)\n",
    "                    if dataset_entity:\n",
    "                        dataset_entity['source_type'] = 'lenta'\n",
    "                        dataset_entity['source_entity'] = f'lenta news'\n",
    "                        lenta_dataset_entities.append(dataset_entity)\n",
    "\n",
    "    # If anything left in the string buffer and the token limit is not\n",
    "    # reached yet - add it to the dataset forcefully.\n",
    "    if lines_candidates:\n",
    "        text_for_entity = ' '.join(lines_candidates)\n",
    "        if text_for_entity:\n",
    "            dataset_entity = make_dataset_entity(text_for_entity)\n",
    "            if dataset_entity:\n",
    "                dataset_entity['source_type'] = 'lenta'\n",
    "                dataset_entity['source_entity'] = f'lenta news'\n",
    "                lenta_dataset_entities.append(dataset_entity)\n",
    "\n",
    "save_entities_to_csv(lenta_dataset_entities, LENTA_SUBSET_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Add a subset from a file to the total dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def add_to_total_dataset(dataset_entities: List[Dict[str, str]], file_path: str):\n",
    "    with open(file_path, 'r') as subset_file:\n",
    "        reader = csv.reader(subset_file)\n",
    "        # Skip the header\n",
    "        next(reader)\n",
    "        for row in reader:\n",
    "            dataset_entities.append({\n",
    "                'source_type': row[0],\n",
    "                'source_entity': row[1],\n",
    "                'text': row[2],\n",
    "                'labels': row[3]\n",
    "            })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make total dataset entities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_dataset_entities: List[Dict[str, str]] = []\n",
    "add_to_total_dataset(total_dataset_entities, SUBTITLES_SUBSET_FILE)\n",
    "add_to_total_dataset(total_dataset_entities, MOSDUMA_SUBSET_FILE)\n",
    "add_to_total_dataset(total_dataset_entities, GOSDUMA_SUBSET_FILE)\n",
    "\n",
    "save_entities_to_csv(total_dataset_entities, TOTAL_SET_FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Method to save text/labels files for subsets of the total dataset (train, dev, test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_text_and_labels_subset(folder_name: str, set_name: str, entities: List[Dict[str, str]]):\n",
    "    with open(os.path.join(folder_name, f\"text_{set_name}.txt\"), \"w\") as text_file:\n",
    "        with open(os.path.join(folder_name, f\"labels_{set_name}.txt\"), \"w\") as labels_file:\n",
    "            for entity in entities:\n",
    "                line_text = entity['text']\n",
    "                line_labels = entity['labels']\n",
    "                text_file.write(line_text + '\\n')\n",
    "                labels_file.write(line_labels + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Splits the result dataset into train/test/dev subsets as separate text and labels files.\n",
    "The data is shuffled before splitting to evenly distribute original\n",
    "corpuses (Gosduma, Mosduma, Subtitles etc.) among subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "total_dataset_entities: List[Dict[str, str]] = []\n",
    "add_to_total_dataset(total_dataset_entities, TOTAL_SET_FILE)\n",
    "random.shuffle(total_dataset_entities)\n",
    "\n",
    "total = len(total_dataset_entities)\n",
    "train, dev = (\n",
    "    int(total * TRAIN_PERCENTS / 100),\n",
    "    int(total * DEV_PERCENTS / 100),\n",
    ")\n",
    "write_text_and_labels_subset(RESULT_DATASET_OUTPUT_DIR, 'train', total_dataset_entities[:train])\n",
    "write_text_and_labels_subset(RESULT_DATASET_OUTPUT_DIR, 'dev', total_dataset_entities[train:(train+dev)])\n",
    "write_text_and_labels_subset(RESULT_DATASET_OUTPUT_DIR, 'test', total_dataset_entities[(train+dev):])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up environmental paths for scripts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# The NeMo scripts should be downloaded separately (and perhaps they are contained in a different folder).\n",
    "nemo_root = os.path.expanduser(os.path.join(os.getcwd(), 'nemo'))\n",
    "config_root = os.path.expanduser(os.path.join(os.getcwd(), 'config'))\n",
    "scripts_root = os.path.expanduser(os.getcwd())\n",
    "dataset_root = os.path.expanduser(RESULT_DATASET_OUTPUT_DIR)\n",
    "\n",
    "os.environ[\"NEMO_ROOT\"] = nemo_root\n",
    "os.environ[\"PUNCTUATOR_CONFIG_ROOT\"] = config_root\n",
    "os.environ[\"PUNCTUATOR_SCRIPTS_ROOT\"] = scripts_root\n",
    "os.environ[\"PUNCTUATOR_DATASET_ROOT\"] = dataset_root\n",
    "\n",
    "print(\"NeMo root:\", nemo_root)\n",
    "print(\"Config root:\", config_root)\n",
    "print(\"Scripts root:\", scripts_root)\n",
    "print(\"Dataset root:\", dataset_root)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "!python $NEMO_ROOT/examples/nlp/token_classification/punctuation_capitalization_train_evaluate.py \\\n",
    "    --config-name=trainer_ru_config.yaml --config-path=$PUNCTUATOR_CONFIG_ROOT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "For inference scripts see \"transcription_inference.ipynb\" nearby."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}