{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23575a1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b15bc7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import orjson\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "893c33d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "podcast_dir = \"/app2/suno/data/podcast/\"\n",
    "aligned_dir = \"/home/victor/neon/sunoData/src/sunodata/podcasts/aligned\"\n",
    "podcast_metas = []\n",
    "\n",
    "# First, read all files\n",
    "all_lines = []\n",
    "for filename in tqdm(os.listdir(aligned_dir)):\n",
    "    if filename.endswith(\".jsonl\"):\n",
    "        filepath = os.path.join(aligned_dir, filename)\n",
    "        with open(filepath) as f:\n",
    "            all_lines.extend(f.readlines())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90349d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Then parse JSON\n",
    "from multiprocessing import Pool\n",
    "\n",
    "\n",
    "def parse_json_line(line):\n",
    "    return orjson.loads(line)\n",
    "\n",
    "\n",
    "with Pool(8) as pool:\n",
    "    podcast_metas = list(\n",
    "        tqdm(pool.imap(parse_json_line, all_lines), total=len(all_lines), desc=\"Parsing JSON\")\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96269871",
   "metadata": {},
   "outputs": [],
   "source": [
    "!head /app2/suno/data/auk_v0/metas_v5_tr.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09be69df",
   "metadata": {},
   "outputs": [],
   "source": [
    "podcast_metas[4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "628defaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "import random\n",
    "\n",
    "# Find an English chunk\n",
    "english_chunks = []\n",
    "for record in podcast_metas:\n",
    "    if record.get(\"chunks\") and record.get(\"detected_language\") == \"english\":\n",
    "        english_chunks.extend(record[\"chunks\"])\n",
    "\n",
    "if english_chunks:\n",
    "    chunk = random.choice(english_chunks)\n",
    "    print(f\"Found English chunk: {chunk['text']}...\")\n",
    "    audio = Audio.from_file(chunk[\"audio_path\"])\n",
    "    audio.play()\n",
    "else:\n",
    "    print(\"No English chunks found\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edc1d165",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Count total chunks across all podcast records\n",
    "total_chunks = 0\n",
    "total_duration = 0\n",
    "records_with_chunks = 0\n",
    "records_without_chunks = 0\n",
    "\n",
    "for record in tqdm(podcast_metas):\n",
    "    if record.get(\"chunks\"):\n",
    "        total_chunks += len(record[\"chunks\"])\n",
    "        records_with_chunks += 1\n",
    "        for chunk in record[\"chunks\"]:\n",
    "            total_duration += chunk[\"duration\"]\n",
    "    else:\n",
    "        records_without_chunks += 1\n",
    "\n",
    "print(f\"Total podcast records: {len(podcast_metas):,}\")\n",
    "print(f\"Records with chunks: {records_with_chunks:,}\")\n",
    "print(f\"Records without chunks: {records_without_chunks:,}\")\n",
    "print(f\"Total chunks: {total_chunks:,}\")\n",
    "print(f\"Average chunks per record (with chunks): {total_chunks / records_with_chunks:.1f}\")\n",
    "print(f\"Average duration per chunk: {total_duration / total_chunks:.1f} seconds\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd71f8f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from collections import Counter\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ea70b9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create character per second distribution plots for each language\n",
    "if podcast_metas:\n",
    "    # Collect characters per second by language\n",
    "    chars_per_second_by_language = {}\n",
    "\n",
    "    for record in tqdm(podcast_metas, desc=\"Collecting chars per second\"):\n",
    "        if record.get(\"chunks\"):\n",
    "            lang = record.get(\"detected_language\", \"unknown\")\n",
    "\n",
    "            for chunk in record[\"chunks\"]:\n",
    "                if \"text\" in chunk and \"duration\" in chunk:\n",
    "                    text_length = len(chunk[\"text\"])\n",
    "                    duration = chunk[\"duration\"]\n",
    "\n",
    "                    if duration > 0:  # Avoid division by zero\n",
    "                        chars_per_second = text_length / duration\n",
    "\n",
    "                        if lang not in chars_per_second_by_language:\n",
    "                            chars_per_second_by_language[lang] = []\n",
    "                        chars_per_second_by_language[lang].append(chars_per_second)\n",
    "\n",
    "    # Get top 20 languages by chunk count for chars per second\n",
    "    language_chunk_counts_cps = [\n",
    "        (lang, len(values)) for lang, values in chars_per_second_by_language.items()\n",
    "    ]\n",
    "    language_chunk_counts_cps.sort(key=lambda x: x[1], reverse=True)\n",
    "    top_20_languages_cps = language_chunk_counts_cps[:20]\n",
    "\n",
    "    # Create subplots\n",
    "    fig, axes = plt.subplots(4, 5, figsize=(20, 25))\n",
    "    fig.suptitle(\"Characters Per Second Distribution by Language (Top 20)\", fontsize=16)\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    for i, (language, chunk_count) in enumerate(top_20_languages_cps):\n",
    "        cps_values = chars_per_second_by_language[language]\n",
    "\n",
    "        # Crop outliers to 95th percentile for better visualization\n",
    "        p95_cps = np.percentile(cps_values, 95)\n",
    "        cropped_cps_values = [x for x in cps_values if x <= p95_cps]\n",
    "\n",
    "        # Create histogram with cropped data\n",
    "        axes[i].hist(cropped_cps_values, bins=30, alpha=0.7, edgecolor=\"black\")\n",
    "        axes[i].set_xlabel(\"Characters Per Second\")\n",
    "        axes[i].set_ylabel(\"Number of Chunks\")\n",
    "        axes[i].set_title(\n",
    "            f\"{language.title()} Chars/Sec Distribution\\n(n={len(cps_values):,} chunks, cropped at 95th)\"\n",
    "        )\n",
    "        axes[i].grid(True, alpha=0.3)\n",
    "\n",
    "        # Add statistics (using original uncropped data)\n",
    "        mean_cps = np.mean(cps_values)\n",
    "        median_cps = np.median(cps_values)\n",
    "        p75_cps = np.percentile(cps_values, 75)\n",
    "        p90_cps = np.percentile(cps_values, 90)\n",
    "\n",
    "        # Only show lines that are within the cropped range\n",
    "        if mean_cps <= p95_cps:\n",
    "            axes[i].axvline(mean_cps, color=\"red\", linestyle=\"--\", label=f\"Mean: {mean_cps:.1f}\")\n",
    "        if median_cps <= p95_cps:\n",
    "            axes[i].axvline(\n",
    "                median_cps, color=\"orange\", linestyle=\"--\", label=f\"Median: {median_cps:.1f}\"\n",
    "            )\n",
    "        if p75_cps <= p95_cps:\n",
    "            axes[i].axvline(p75_cps, color=\"green\", linestyle=\":\", label=f\"75th: {p75_cps:.1f}\")\n",
    "        if p90_cps <= p95_cps:\n",
    "            axes[i].axvline(p90_cps, color=\"blue\", linestyle=\":\", label=f\"90th: {p90_cps:.1f}\")\n",
    "\n",
    "        axes[i].legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Print statistics for each language\n",
    "    print(f\"\\nCharacters Per Second Statistics by Language:\")\n",
    "    for language, chunk_count in top_20_languages_cps:\n",
    "        cps_values = chars_per_second_by_language[language]\n",
    "        print(f\"\\n{language.title()}:\")\n",
    "        print(f\"  Count: {len(cps_values):,}\")\n",
    "        print(f\"  Mean: {np.mean(cps_values):.2f}\")\n",
    "        print(f\"  Median: {np.median(cps_values):.2f}\")\n",
    "        print(f\"  75th percentile: {np.percentile(cps_values, 75):.2f}\")\n",
    "        print(f\"  90th percentile: {np.percentile(cps_values, 90):.2f}\")\n",
    "        print(f\"  95th percentile: {np.percentile(cps_values, 95):.2f}\")\n",
    "        print(f\"  Min: {min(cps_values):.2f}\")\n",
    "        print(f\"  Max: {max(cps_values):.2f}\")\n",
    "        print(f\"  Std: {np.std(cps_values):.2f}\")\n",
    "else:\n",
    "    print(\"No character per second values found in chunks\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02ef2673",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter chunks to keep only 25-75 percentiles within each language by CPS\n",
    "filtered_podcast_metas = []\n",
    "\n",
    "for record in tqdm(podcast_metas, desc=\"Filtering by CPS percentiles\"):\n",
    "    if not record.get(\"chunks\") or not record.get(\"detected_language\"):\n",
    "        continue\n",
    "\n",
    "    language = record[\"detected_language\"]\n",
    "\n",
    "    # Calculate CPS for all chunks in this record\n",
    "    chunk_cps_values = []\n",
    "    for chunk in record[\"chunks\"]:\n",
    "        if \"text\" in chunk and \"duration\" in chunk and chunk[\"duration\"] > 0:\n",
    "            cps = len(chunk[\"text\"]) / chunk[\"duration\"]\n",
    "            chunk_cps_values.append(cps)\n",
    "        else:\n",
    "            chunk_cps_values.append(None)\n",
    "\n",
    "    # Get language-specific CPS percentiles\n",
    "    if language in chars_per_second_by_language:\n",
    "        lang_cps_values = chars_per_second_by_language[language]\n",
    "        p25_cps = np.percentile(lang_cps_values, 25)\n",
    "        p75_cps = np.percentile(lang_cps_values, 75)\n",
    "\n",
    "        # Filter chunks that fall within 25-75 percentile range\n",
    "        filtered_chunks = []\n",
    "        for chunk, cps in zip(record[\"chunks\"], chunk_cps_values):\n",
    "            if cps is not None and p25_cps <= cps <= p75_cps:\n",
    "                filtered_chunks.append(chunk)\n",
    "\n",
    "        # Only keep records that have at least one chunk after filtering\n",
    "        if filtered_chunks:\n",
    "            filtered_record = record.copy()\n",
    "            filtered_record[\"chunks\"] = filtered_chunks\n",
    "            filtered_podcast_metas.append(filtered_record)\n",
    "\n",
    "print(f\"Original records: {len(podcast_metas):,}\")\n",
    "print(f\"Filtered records: {len(filtered_podcast_metas):,}\")\n",
    "\n",
    "# Count total chunks before and after filtering\n",
    "original_chunk_count = sum(len(record.get(\"chunks\", [])) for record in podcast_metas)\n",
    "filtered_chunk_count = sum(len(record.get(\"chunks\", [])) for record in filtered_podcast_metas)\n",
    "\n",
    "print(f\"Original chunks: {original_chunk_count:,}\")\n",
    "print(f\"Filtered chunks: {filtered_chunk_count:,}\")\n",
    "print(f\"Retention rate: {filtered_chunk_count / original_chunk_count * 100:.1f}%\")\n",
    "\n",
    "# Update podcast_metas to use filtered data\n",
    "podcast_metas = filtered_podcast_metas\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed3eee5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Collect CER values by language\n",
    "cer_by_language = {}\n",
    "for record in tqdm(podcast_metas):\n",
    "    if record.get(\"chunks\") and record.get(\"detected_language\"):\n",
    "        language = record[\"detected_language\"]\n",
    "        if language not in cer_by_language:\n",
    "            cer_by_language[language] = []\n",
    "\n",
    "        for chunk in record[\"chunks\"]:\n",
    "            if \"hoot_cer\" in chunk:\n",
    "                cer_by_language[language].append(chunk[\"hoot_cer\"])\n",
    "\n",
    "# Find the 20 most popular languages by number of chunks\n",
    "language_counts = {lang: len(cer_values) for lang, cer_values in cer_by_language.items()}\n",
    "top_20_languages = sorted(language_counts.items(), key=lambda x: x[1], reverse=True)[:20]\n",
    "\n",
    "print(f\"Top 20 languages by chunk count:\")\n",
    "for lang, count in top_20_languages:\n",
    "    print(f\"  {lang}: {count:,} chunks\")\n",
    "\n",
    "# Create subplots for the top 20 languages\n",
    "if top_20_languages:\n",
    "    fig, axes = plt.subplots(5, 4, figsize=(20, 25))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    for i, (language, chunk_count) in enumerate(top_20_languages):\n",
    "        cer_values = cer_by_language[language]\n",
    "\n",
    "        # Create histogram\n",
    "        axes[i].hist(cer_values, bins=30, alpha=0.7, edgecolor=\"black\")\n",
    "        axes[i].set_xlabel(\"Character Error Rate (CER)\")\n",
    "        axes[i].set_ylabel(\"Number of Chunks\")\n",
    "        axes[i].set_title(f\"{language.title()} CER Distribution\\n(n={len(cer_values):,} chunks)\")\n",
    "        axes[i].grid(True, alpha=0.3)\n",
    "\n",
    "        # Add statistics\n",
    "        mean_cer = np.mean(cer_values)\n",
    "        median_cer = np.median(cer_values)\n",
    "        p75_cer = np.percentile(cer_values, 75)\n",
    "        p90_cer = np.percentile(cer_values, 90)\n",
    "        p95_cer = np.percentile(cer_values, 95)\n",
    "\n",
    "        axes[i].axvline(mean_cer, color=\"red\", linestyle=\"--\", label=f\"Mean: {mean_cer:.3f}\")\n",
    "        axes[i].axvline(median_cer, color=\"orange\", linestyle=\"--\", label=f\"Median: {median_cer:.3f}\")\n",
    "        axes[i].axvline(p75_cer, color=\"green\", linestyle=\":\", label=f\"75th: {p75_cer:.3f}\")\n",
    "        axes[i].axvline(p90_cer, color=\"blue\", linestyle=\":\", label=f\"90th: {p90_cer:.3f}\")\n",
    "        axes[i].axvline(p95_cer, color=\"purple\", linestyle=\":\", label=f\"95th: {p95_cer:.3f}\")\n",
    "        axes[i].legend()\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    # Print statistics for each language\n",
    "    print(f\"\\nCER Statistics by Language:\")\n",
    "    for language, chunk_count in top_20_languages:\n",
    "        cer_values = cer_by_language[language]\n",
    "        print(f\"\\n{language.title()}:\")\n",
    "        print(f\"  Count: {len(cer_values):,}\")\n",
    "        print(f\"  Mean: {np.mean(cer_values):.4f}\")\n",
    "        print(f\"  Median: {np.median(cer_values):.4f}\")\n",
    "        print(f\"  75th percentile: {np.percentile(cer_values, 75):.4f}\")\n",
    "        print(f\"  90th percentile: {np.percentile(cer_values, 90):.4f}\")\n",
    "        print(f\"  95th percentile: {np.percentile(cer_values, 95):.4f}\")\n",
    "        print(f\"  Min: {min(cer_values):.4f}\")\n",
    "        print(f\"  Max: {max(cer_values):.4f}\")\n",
    "        print(f\"  Std: {np.std(cer_values):.4f}\")\n",
    "else:\n",
    "    print(\"No CER values found in chunks\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce4d1110",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter chunks to only include those under the mean CER for each language\n",
    "if cer_values:\n",
    "    # Calculate mean CER for each language\n",
    "    mean_cer_by_language = {}\n",
    "    for language, chunk_count in top_20_languages:\n",
    "        if language in cer_by_language:\n",
    "            mean_cer_by_language[language] = np.mean(cer_by_language[language])\n",
    "\n",
    "    # Filter podcast_metas to only include chunks under mean CER\n",
    "    filtered_podcast_metas = []\n",
    "    total_chunks_before = 0\n",
    "    total_chunks_after = 0\n",
    "\n",
    "    for record in tqdm(podcast_metas, desc=\"Filtering chunks by mean CER\"):\n",
    "        if record.get(\"chunks\"):\n",
    "            lang = record.get(\"detected_language\", \"unknown\")\n",
    "            filtered_chunks = []\n",
    "\n",
    "            total_chunks_before += len(record[\"chunks\"])\n",
    "\n",
    "            if lang in mean_cer_by_language:\n",
    "                mean_cer = mean_cer_by_language[lang]\n",
    "\n",
    "                for chunk in record[\"chunks\"]:\n",
    "                    if \"hoot_cer\" in chunk and chunk[\"hoot_cer\"] < mean_cer:\n",
    "                        filtered_chunks.append(chunk)\n",
    "\n",
    "            # Only keep records that have at least one chunk after filtering\n",
    "            if filtered_chunks:\n",
    "                filtered_record = record.copy()\n",
    "                filtered_record[\"chunks\"] = filtered_chunks\n",
    "                filtered_podcast_metas.append(filtered_record)\n",
    "                total_chunks_after += len(filtered_chunks)\n",
    "\n",
    "    print(f\"\\nFiltering results:\")\n",
    "    print(f\"Total chunks before filtering: {total_chunks_before:,}\")\n",
    "    print(f\"Total chunks after filtering: {total_chunks_after:,}\")\n",
    "    print(f\"Chunks removed: {total_chunks_before - total_chunks_after:,}\")\n",
    "    print(f\"Retention rate: {total_chunks_after / total_chunks_before * 100:.1f}%\")\n",
    "    print(f\"Records before filtering: {len(podcast_metas):,}\")\n",
    "    print(f\"Records after filtering: {len(filtered_podcast_metas):,}\")\n",
    "\n",
    "    # Update podcast_metas to use filtered version\n",
    "    podcast_metas = filtered_podcast_metas\n",
    "else:\n",
    "    print(\"No CER values found for filtering\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11168ad3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listen to 3 random chunks per language\n",
    "import random\n",
    "\n",
    "# Group chunks by language\n",
    "chunks_by_language = {}\n",
    "for record in podcast_metas:\n",
    "    lang = record.get(\"detected_language\", \"unknown\")\n",
    "    if lang not in chunks_by_language:\n",
    "        chunks_by_language[lang] = []\n",
    "\n",
    "    for chunk in record[\"chunks\"]:\n",
    "        chunks_by_language[lang].append(\n",
    "            {\n",
    "                \"audio_path\": chunk[\"audio_path\"],\n",
    "                \"text\": chunk[\"text\"],\n",
    "                \"duration\": chunk[\"duration\"],\n",
    "                \"hoot_cer\": chunk.get(\"hoot_cer\", \"N/A\"),\n",
    "                \"record_id\": record[\"id\"],\n",
    "            }\n",
    "        )\n",
    "\n",
    "# Sample and display 3 random chunks per language\n",
    "for lang, chunks in chunks_by_language.items():\n",
    "    if len(chunks) == 0:\n",
    "        continue\n",
    "\n",
    "    print(f\"\\n{'=' * 50}\")\n",
    "    print(f\"Language: {lang} ({len(chunks)} total chunks)\")\n",
    "    print(f\"{'=' * 50}\")\n",
    "\n",
    "    # Sample up to 3 random chunks\n",
    "    sample_size = min(1, len(chunks))\n",
    "    sampled_chunks = random.sample(chunks, sample_size)\n",
    "\n",
    "    for i, chunk in enumerate(sampled_chunks, 1):\n",
    "        print(f\"\\nChunk {i}/3:\")\n",
    "        print(f\"Duration: {chunk['duration']:.1f}s\")\n",
    "        print(f\"CER: {chunk['hoot_cer']}\")\n",
    "        print(f\"Text preview: {chunk['text'][:200]}...\")\n",
    "\n",
    "        # Display audio player\n",
    "        audio = Audio.from_file(chunk[\"audio_path\"])\n",
    "        audio.play()\n",
    "        print(\"-\" * 30)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4dd66f9",
   "metadata": {},
   "source": [
    "## Other stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dd3b3e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_podcast_record(record):\n",
    "    \"\"\"Transform podcast format to SunoGPT metadata format\"\"\"\n",
    "\n",
    "    # Extract basic info\n",
    "    track_id = str(record[\"id\"])  # Convert to string for consistency\n",
    "\n",
    "    # Build tags from various fields\n",
    "    tags = []\n",
    "\n",
    "    # Add detected language\n",
    "    if record.get(\"detected_language\"):\n",
    "        tags.append(f\"language: {record['detected_language']}\")\n",
    "\n",
    "    # Add podcast-specific tags\n",
    "    tags.append(\"podcast\")\n",
    "    tags.append(\"speech\")\n",
    "\n",
    "    suno_records = []\n",
    "\n",
    "    # dedup chunks by chunk_index\n",
    "    seen_chunks = set()\n",
    "\n",
    "    for chunk in record[\"chunks\"]:\n",
    "        if chunk[\"chunk_index\"] in seen_chunks:\n",
    "            continue\n",
    "        seen_chunks.add(chunk[\"chunk_index\"])\n",
    "\n",
    "        text_aligned = chunk[\"hoot_alignment\"]\n",
    "        text_aligned = [(w[\"start_s\"], w[\"end_s\"], w[\"word\"]) for w in text_aligned]\n",
    "\n",
    "        # Build the SunoGPT metadata record\n",
    "        assert os.path.exists(chunk[\"audio_path\"])\n",
    "        suno_record = {\n",
    "            \"id\": f\"podcast_{track_id}_{chunk['chunk_index']}\",  # Prefix to avoid ID conflicts\n",
    "            \"audio_type\": \"speech\",\n",
    "            \"local_filepath\": chunk[\"audio_path\"],\n",
    "            \"duration_s\": float(chunk[\"duration\"]),\n",
    "            \"tags\": tags,\n",
    "            \"lang\": record[\"detected_language\"],\n",
    "            \"text\": chunk[\"text\"],\n",
    "            \"source\": \"podcast\",\n",
    "            \"artist_ids\": [f\"podcast_{record['id']}\"],\n",
    "            \"hoot_cer\": chunk[\"hoot_cer\"],\n",
    "            \"text_aligned\": text_aligned,\n",
    "        }\n",
    "\n",
    "        suno_records.append(suno_record)\n",
    "\n",
    "    return suno_records\n",
    "\n",
    "\n",
    "# Test the transformation with a sample record\n",
    "print(\"Testing transformation function...\")\n",
    "transform_podcast_record(podcast_metas[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a2392f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process all extreme music records\n",
    "print(f\"Processing {len(podcast_metas)} podcast records...\")\n",
    "\n",
    "from multiprocessing import Pool\n",
    "import multiprocessing as mp\n",
    "\n",
    "\n",
    "def process_record_wrapper(record):\n",
    "    try:\n",
    "        return transform_podcast_record(record)\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing record: {e}\")\n",
    "        return None\n",
    "\n",
    "\n",
    "with Pool(processes=8) as pool:\n",
    "    results = list(tqdm(pool.imap(process_record_wrapper, podcast_metas), total=len(podcast_metas)))\n",
    "\n",
    "suno_records = [record for record in results if record is not None]\n",
    "print(f\"Successfully processed {len(suno_records)} records\")\n",
    "# flatten list\n",
    "suno_records = [item for sublist in suno_records for item in sublist]\n",
    "print(f\"Flattened list to {len(suno_records)} records\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8498aea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "# Split suno_records into train and val\n",
    "random.seed(42)\n",
    "random.shuffle(suno_records)\n",
    "\n",
    "split_ratio = 0.95\n",
    "split_index = int(len(suno_records) * split_ratio)\n",
    "\n",
    "suno_records_train = suno_records[:split_index]\n",
    "suno_records_val = suno_records[split_index:]\n",
    "\n",
    "print(\n",
    "    f\"Split {len(suno_records)} records into {len(suno_records_train)} train and {len(suno_records_val)} val\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c7093e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp /app2/suno/data/auk_v0/metas_v5_val.jsonl /tmp/metas_v6_val.jsonl\n",
    "!cp /app2/suno/data/auk_v0/metas_v5_tr.jsonl /tmp/metas_v6_tr.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fb1ec83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Append extreme music records to existing v6 files\n",
    "import json\n",
    "\n",
    "# Append to train file\n",
    "train_file = \"/tmp/metas_v6_tr.jsonl\"\n",
    "with open(train_file, \"a\") as f:\n",
    "    for record in suno_records_train:\n",
    "        f.write(json.dumps(record) + \"\\n\")\n",
    "\n",
    "# Append to val file\n",
    "val_file = \"/tmp/metas_v6_val.jsonl\"\n",
    "with open(val_file, \"a\") as f:\n",
    "    for record in suno_records_val:\n",
    "        f.write(json.dumps(record) + \"\\n\")\n",
    "\n",
    "print(f\"Appended {len(suno_records_train)} extreme records to train file\")\n",
    "print(f\"Appended {len(suno_records_val)} extreme records to val file\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eadaa562",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sudo cp /tmp/metas_v6_tr.jsonl /app2/suno/data/auk_v0/metas_v6_tr.jsonl\n",
    "!sudo cp /tmp/metas_v6_val.jsonl /app2/suno/data/auk_v0/metas_v6_val.jsonl"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}