{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de57d37b-73e7-40aa-a922-7c4a9fbc8085",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Обработка тренировочного датасета...\n",
      "Загружаем список файлов из ngn/train_big.tsv\n",
      "Проверяем размеры 1120 файлов (ищем файлы < 1.0MB)...\n",
      "Файл слишком большой: ngn/clips/Audio_Ngen_2019_2020_2021/02012020/995_02012020.wav (2.06MB > 1.0MB)\n",
      "Подходящий файл найден: ngn/clips/Audio_Ngen_2019_2020_2021/26122019/110.wav, размер: 0.79MB\n",
      "Подходящий файл найден: ngn/clips/Audio_Ngen_2019_2020_2021/27122019/321untitled.wav, размер: 0.43MB\n",
      "Подходящий файл найден: ngn/clips/Audio_Ngen_2019_2020_2021/06122020/1448untitled.wav, размер: 0.55MB\n",
      "Подходящий файл найден: ngn/clips/Audio_Ngen_2019_2020_2021/10012020/1866_10012020.wav, размер: 0.22MB\n",
      "Подходящий файл найден: ngn/clips/Audio_Ngen_2019_2020_2021/12012020/1986untitled.wav, размер: 0.32MB\n",
      "Файл слишком большой: ngn/clips/Audio_Ngen_2019_2020_2021/30122019/658untitled.wav (1.26MB > 1.0MB)\n",
      "Файл слишком большой: ngn/clips/Audio_Ngen_2019_2020_2021/03012020/1064untitled.wav (1.52MB > 1.0MB)\n",
      "Файл слишком большой: ngn/clips/Audio_Ngen_2019_2020_2021/03012020/1098untitled.wav (1.09MB > 1.0MB)\n",
      "Файл слишком большой: ngn/clips/Audio_Ngen_2019_2020_2021/26122019/116.wav (1.17MB > 1.0MB)\n",
      "Проверено 100/1120 файлов...\n",
      "Проверено 200/1120 файлов...\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from google.cloud import storage\n",
    "import pandas as pd\n",
    "from typing import Tuple, List, Dict\n",
    "import io\n",
    "\n",
    "def check_blob_exists_and_size(bucket, possible_paths: List[str]) -> Tuple[bool, str, float]:\n",
    "    \"\"\"\n",
    "    Проверяет существование файла по всем возможным путям и его размер\n",
    "    Returns:\n",
    "        Tuple[bool, str, float]: (найден ли файл, путь к файлу, размер в MB)\n",
    "    \"\"\"\n",
    "    for path in possible_paths:\n",
    "        blob = bucket.blob(path)\n",
    "        try:\n",
    "            if blob.exists():\n",
    "                blob.reload()  # Загружаем метаданные\n",
    "                size_mb = blob.size / (1024 * 1024)\n",
    "                return True, path, size_mb\n",
    "        except Exception as e:\n",
    "            print(f\"Ошибка при проверке файла {path}: {str(e)}\")\n",
    "            continue\n",
    "    return False, \"\", 0.0\n",
    "\n",
    "def generate_possible_paths(original_path: str, clips_prefix: str) -> List[str]:\n",
    "    \"\"\"\n",
    "    Генерирует все возможные варианты пути к файлу с учетом префикса\n",
    "    \"\"\"\n",
    "    parts = original_path.split('/')\n",
    "    if len(parts) < 2:\n",
    "        return []\n",
    "    current_folder = parts[-2]  # получаем название текущей папки\n",
    "    filename = parts[-1]  # получаем имя файла с расширением\n",
    "    base_dir = '/'.join(parts[:-1])  # получаем путь к директории\n",
    "\n",
    "    # Разделяем имя файла на имя и расширение\n",
    "    name, extension = os.path.splitext(filename)\n",
    "    \n",
    "    # Разделяем имя по знаку подчеркивания и берем левую часть\n",
    "    base_name = name.split('_')[0]\n",
    "    \n",
    "    # Формируем все возможные варианты имен файлов\n",
    "    possible_names = [\n",
    "        f\"{base_name}{extension}\",  # просто имя с расширением\n",
    "        f\"{base_name}_{current_folder}{extension}\",  # имя + текущая папка\n",
    "        f\"{base_name}untitled{extension}\"  # имя + untitled\n",
    "    ]\n",
    "    \n",
    "    # Формируем полные пути с префиксом\n",
    "    possible_paths = [f\"{clips_prefix}/{base_dir}/{name}\" for name in possible_names]\n",
    "    \n",
    "    return list(set(possible_paths))  # удаляем возможные дубликаты\n",
    "\n",
    "def check_all_files_sizes(bucket_name: str, tsv_path: str, clips_prefix: str, max_size_mb: float = 1.0) -> Dict[str, Tuple[str, float]]:\n",
    "    \"\"\"\n",
    "    Проверяет размеры всех файлов и возвращает словарь с информацией о файлах меньше max_size_mb\n",
    "    \"\"\"\n",
    "    storage_client = storage.Client()\n",
    "    bucket = storage_client.bucket(bucket_name)\n",
    "    \n",
    "    print(f\"Загружаем список файлов из {tsv_path}\")\n",
    "    blob = bucket.blob(tsv_path)\n",
    "    content = blob.download_as_string()\n",
    "    df = pd.read_csv(io.BytesIO(content), sep='\\t')\n",
    "    \n",
    "    valid_files = {}\n",
    "    missing_files = []\n",
    "    oversized_files = []\n",
    "    \n",
    "    total_files = len(df)\n",
    "    print(f\"Проверяем размеры {total_files} файлов (ищем файлы < {max_size_mb}MB)...\")\n",
    "    \n",
    "    # Статистика по размерам\n",
    "    size_stats = {\n",
    "        '0-0.5MB': 0,\n",
    "        '0.5-1MB': 0,\n",
    "        '1-1.5MB': 0,\n",
    "        '1.5-2MB': 0,\n",
    "        '>2MB': 0\n",
    "    }\n",
    "    \n",
    "    # Проверяем каждый файл\n",
    "    for idx, row in df.iterrows():\n",
    "        if 'path' not in row or pd.isna(row['path']):\n",
    "            continue\n",
    "            \n",
    "        original_path = row['path'].strip()\n",
    "        possible_paths = generate_possible_paths(original_path, clips_prefix)\n",
    "        exists, real_path, size_mb = check_blob_exists_and_size(bucket, possible_paths)\n",
    "        \n",
    "        if exists:\n",
    "            # Обновляем статистику по размерам\n",
    "            if size_mb > 2:\n",
    "                size_stats['>2MB'] += 1\n",
    "            elif size_mb > 1.5:\n",
    "                size_stats['1.5-2MB'] += 1\n",
    "            elif size_mb > 1:\n",
    "                size_stats['1-1.5MB'] += 1\n",
    "            elif size_mb > 0.5:\n",
    "                size_stats['0.5-1MB'] += 1\n",
    "            else:\n",
    "                size_stats['0-0.5MB'] += 1\n",
    "            \n",
    "            if size_mb < max_size_mb:\n",
    "                valid_files[original_path] = (real_path, size_mb)\n",
    "                if len(valid_files) <= 5:\n",
    "                    print(f\"Подходящий файл найден: {real_path}, размер: {size_mb:.2f}MB\")\n",
    "            else:\n",
    "                oversized_files.append((original_path, size_mb))\n",
    "                if len(oversized_files) <= 5:\n",
    "                    print(f\"Файл слишком большой: {real_path} ({size_mb:.2f}MB > {max_size_mb}MB)\")\n",
    "        else:\n",
    "            missing_files.append(original_path)\n",
    "            \n",
    "        if (idx + 1) % 100 == 0:\n",
    "            print(f\"Проверено {idx + 1}/{total_files} файлов...\")\n",
    "    \n",
    "    # Выводим статистику\n",
    "    print(\"\\nРезультаты проверки размеров:\")\n",
    "    print(f\"Всего файлов: {total_files}\")\n",
    "    print(f\"Файлов не найдено: {len(missing_files)}\")\n",
    "    print(f\"\\nРаспределение по размерам:\")\n",
    "    for size_range, count in size_stats.items():\n",
    "        print(f\"{size_range}: {count} файлов\")\n",
    "    \n",
    "    return valid_files, missing_files, oversized_files\n",
    "\n",
    "def create_filtered_dataset(bucket_name: str, tsv_path: str, clips_prefix: str, max_size_mb: float = 1.0) -> pd.DataFrame:\n",
    "    \"\"\"\n",
    "    Создает датафрейм только с файлами подходящего размера\n",
    "    \"\"\"\n",
    "    # Проверяем размеры всех файлов\n",
    "    valid_files, missing_files, oversized_files = check_all_files_sizes(\n",
    "        bucket_name, tsv_path, clips_prefix, max_size_mb\n",
    "    )\n",
    "    \n",
    "    # Загружаем исходный датафрейм\n",
    "    storage_client = storage.Client()\n",
    "    bucket = storage_client.bucket(bucket_name)\n",
    "    blob = bucket.blob(tsv_path)\n",
    "    content = blob.download_as_string()\n",
    "    df = pd.read_csv(io.BytesIO(content), sep='\\t')\n",
    "    \n",
    "    # Очищаем датафрейм\n",
    "    df = df.dropna(subset=['path', 'sentence_normalized'])\n",
    "    df = df[\n",
    "        (df['path'].str.strip() != '') & \n",
    "        (df['sentence_normalized'].str.strip() != '')\n",
    "    ]\n",
    "    \n",
    "    # Фильтруем датафрейм\n",
    "    filtered_df = df[df['path'].isin(valid_files.keys())].copy()\n",
    "    \n",
    "    # Обновляем пути и создаем аудио колонку\n",
    "    filtered_df['path'] = filtered_df['path'].apply(lambda x: valid_files[x][0])\n",
    "    filtered_df['audio'] = filtered_df['path'].apply(lambda x: f\"gs://{bucket_name}/{x}\")\n",
    "    \n",
    "    # Создаем тройную транскрипцию\n",
    "    filtered_df['sentence_normalized'] = filtered_df['sentence_normalized'].apply(\n",
    "        lambda x: ' '.join([x.strip()] * 3)\n",
    "    )\n",
    "    \n",
    "    print(f\"\\nСоздан отфильтрованный датасет с {len(filtered_df)} записями\")\n",
    "    return filtered_df\n",
    "\n",
    "def process_datasets(bucket_name: str, train_tsv: str, test_tsv: str, clips_prefix: str, max_size_mb: float = 1.0):\n",
    "    \"\"\"\n",
    "    Обрабатывает тренировочный и тестовый датасеты\n",
    "    \"\"\"\n",
    "    print(\"Обработка тренировочного датасета...\")\n",
    "    filtered_train_df = create_filtered_dataset(\n",
    "        bucket_name=bucket_name,\n",
    "        tsv_path=train_tsv,\n",
    "        clips_prefix=clips_prefix,\n",
    "        max_size_mb=max_size_mb\n",
    "    )\n",
    "    \n",
    "    print(\"\\nОбработка тестового датасета...\")\n",
    "    filtered_test_df = create_filtered_dataset(\n",
    "        bucket_name=bucket_name,\n",
    "        tsv_path=test_tsv,\n",
    "        clips_prefix=clips_prefix,\n",
    "        max_size_mb=max_size_mb\n",
    "    )\n",
    "    \n",
    "    # Сохраняем результаты\n",
    "    filtered_train_df.to_csv('1filtered_train_dataset.tsv', sep='\\t', index=False)\n",
    "    filtered_test_df.to_csv('filtered_test_dataset.tsv', sep='\\t', index=False)\n",
    "    \n",
    "    return filtered_train_df, filtered_test_df\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # Обрабатываем оба датасета\n",
    "    filtered_train_df, filtered_test_df = process_datasets(\n",
    "        bucket_name='ngen_model_fine_tuned',\n",
    "        train_tsv='ngn/train_big.tsv',\n",
    "        test_tsv='ngn/test_big.tsv',\n",
    "        clips_prefix='ngn/clips',\n",
    "        max_size_mb=1.0  # ограничение в 1MB\n",
    "    )\n",
    "    \n",
    "#     # Создаем датасеты для huggingface\n",
    "#     from datasets import Dataset\n",
    "#     from datasets.features import Audio\n",
    "    \n",
    "#     train_dataset = Dataset.from_pandas(filtered_train_df)\n",
    "#     test_dataset = Dataset.from_pandas(filtered_test_df)\n",
    "#     print(\"Созданы датасеты\")\n",
    "    \n",
    "#     train_dataset = train_dataset.cast_column(\"audio\", Audio(sampling_rate=16000))\n",
    "#     test_dataset = test_dataset.cast_column(\"audio\", Audio(sampling_rate=16000))\n",
    "#     print(\"Колонки аудио преобразованы\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89853d1c-a7db-498b-bd45-bc6e72b86b4e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m126",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m126"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}