{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "0", "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "from suno_utils.utils.text import read_jsonl\n", "from suno_utils.utils.s3 import read_from_s3, upload_s3_files\n", "\n", "import pandas as pd\n", "import numpy as np\n", "from collections import Counter\n", "from sklearn.model_selection import train_test_split\n", "import warnings\n", "from tqdm import tqdm\n", "import pandas as pd\n", "import numpy as np\n", "from collections import Counter\n", "from sklearn.model_selection import train_test_split\n", "import warnings\n", "from multiprocessing import Pool, cpu_count\n", "from functools import partial\n", "import ast\n", "from tqdm import tqdm\n", "import tempfile" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "raw_metas = read_from_s3(\"s3://suno-data/datasets/harvest/pond5_sfx/pond5_metas_v0.jsonl\", read_f=read_jsonl)\n", "print(f\"{len(raw_metas):,} tracks with {sum([m['duration_s'] for m in raw_metas])/60/60:,.1f}h total\")\n", "df_raw = pd.DataFrame.from_dict(raw_metas)\n", "print(df_raw.columns)\n", "df_raw.head()" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "### Splice" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "def process_pack_batch(pack_batch_with_id):\n", " \"\"\"Process a batch of packs to extract dominant tags\"\"\"\n", " batch_id, pack_batch, df_subset = pack_batch_with_id\n", " results = []\n", " \n", " # Create progress bar for this batch\n", " pbar = tqdm(pack_batch, desc=f\"Batch {batch_id}\", leave=False)\n", " \n", " for pack in pbar:\n", " pack_data = df_subset[df_subset['parent_slug'] == pack]\n", " \n", " # Fast tag processing\n", " all_tags = []\n", " for tag_list in pack_data['processed_tags'].dropna():\n", " if isinstance(tag_list, list):\n", " all_tags.extend(tag_list)\n", " elif isinstance(tag_list, str) and tag_list.startswith('['):\n", " # Handle string representation of lists\n", " try:\n", " all_tags.extend(ast.literal_eval(tag_list))\n", " except:\n", " pass\n", " \n", " dominant_tag = Counter(all_tags).most_common(1)[0][0] if all_tags else None\n", " results.append((pack, dominant_tag))\n", " \n", " pbar.close()\n", " return results\n", "\n", "def create_stratified_pack_split_fast(df, test_size=0.2, random_state=42, n_processes=None):\n", " \"\"\"\n", " Ultra-fast stratified split using vectorized operations and multiprocessing.\n", " \"\"\"\n", " if n_processes is None:\n", " n_processes = min(cpu_count(), 8) # Don't use too many cores\n", " \n", " print(f\"Using {n_processes} processes\")\n", " print(f\"Total sample packs: {df['parent_slug'].nunique()}\")\n", " print(f\"Total samples: {len(df)}\")\n", " \n", " # Fast pack-level aggregation using vectorized operations\n", " print(\"Computing pack statistics...\")\n", " pack_stats = df.groupby('parent_slug').agg({\n", " 'duration_s': ['mean', 'std', 'count'],\n", " 'bpm': lambda x: pd.to_numeric(x, errors='coerce').mean(),\n", " 'key': lambda x: x.mode().iloc[0] if x.notna().any() else None\n", " })\n", " \n", " # Flatten column names\n", " pack_stats.columns = ['avg_duration', 'duration_std', 'sample_count', 'avg_bpm', 'most_common_key']\n", " pack_stats = pack_stats.reset_index()\n", " \n", " # Parallel tag processing\n", " print(\"Processing tags in parallel...\")\n", " unique_packs = pack_stats['parent_slug'].tolist()\n", " \n", " # Split packs into batches for parallel processing\n", " batch_size = max(1, len(unique_packs) // (n_processes * 4)) # 4 batches per process\n", " pack_batches = [unique_packs[i:i + batch_size] for i in range(0, len(unique_packs), batch_size)]\n", " \n", " # Create subset of df with only needed columns for faster processing\n", " print(\"Preparing data for parallel processing...\")\n", " df_subset = df[['parent_slug', 'processed_tags']].copy()\n", " \n", " # Add batch IDs and df_subset to each batch\n", " pack_batches_with_ids = [(i, batch, df_subset) for i, batch in enumerate(pack_batches)]\n", " \n", " # Process in parallel with overall progress bar\n", " print(f\"Processing {len(pack_batches)} batches across {n_processes} processes...\")\n", " \n", " with Pool(n_processes) as pool:\n", " # Use imap for progress tracking\n", " batch_results = list(tqdm(\n", " pool.imap(process_pack_batch, pack_batches_with_ids),\n", " total=len(pack_batches),\n", " desc=\"Processing batches\"\n", " ))\n", " \n", " # Flatten results\n", " tag_results = {}\n", " for batch_result in batch_results:\n", " for pack, tag in batch_result:\n", " tag_results[pack] = tag\n", " \n", " # Add tags to pack_stats\n", " pack_stats['dominant_tag'] = pack_stats['parent_slug'].map(tag_results)\n", " \n", " print(\"Creating stratification bins...\")\n", " # Simple binning for stratification\n", " pack_stats['duration_bin'] = pd.cut(pack_stats['avg_duration'], bins=3, labels=['short', 'medium', 'long'])\n", " \n", " # Create simple stratification key (just tag + duration for speed)\n", " pack_stats['strat_key'] = (pack_stats['dominant_tag'].fillna('unknown').astype(str) + '_' + \n", " pack_stats['duration_bin'].astype(str))\n", " \n", " print(\"Splitting packs...\")\n", " # Split packs\n", " try:\n", " # Only stratify if we have multiple packs per stratum\n", " strat_counts = pack_stats['strat_key'].value_counts()\n", " if (strat_counts >= 2).sum() >= len(strat_counts) * 0.3:\n", " train_packs, test_packs = train_test_split(\n", " pack_stats['parent_slug'], \n", " test_size=test_size,\n", " stratify=pack_stats['strat_key'],\n", " random_state=random_state\n", " )\n", " else:\n", " raise ValueError(\"Too few samples per stratum\")\n", " except:\n", " print(\"Using simple random split by packs\")\n", " train_packs, test_packs = train_test_split(\n", " pack_stats['parent_slug'], \n", " test_size=test_size,\n", " random_state=random_state\n", " )\n", " \n", " print(\"Creating final splits...\")\n", " # Use vectorized operations for final split\n", " train_pack_set = set(train_packs) # Convert to set for faster lookup\n", " \n", " # Add progress bar for the final split operation\n", " tqdm.pandas(desc=\"Creating train/test masks\")\n", " train_mask = df['parent_slug'].progress_apply(lambda x: x in train_pack_set)\n", " \n", " train_df = df[train_mask].copy()\n", " test_df = df[~train_mask].copy()\n", " \n", " print(f\"\\nSplit Results:\")\n", " print(f\"Train packs: {len(train_packs)} ({len(train_packs)/len(pack_stats)*100:.1f}%)\")\n", " print(f\"Test packs: {len(test_packs)} ({len(test_packs)/len(pack_stats)*100:.1f}%)\")\n", " print(f\"Train samples: {len(train_df)} ({len(train_df)/len(df)*100:.1f}%)\")\n", " print(f\"Test samples: {len(test_df)} ({len(test_df)/len(df)*100:.1f}%)\")\n", " \n", " return train_df, test_df\n", "\n", "def analyze_split_quality_fast(train_df, test_df):\n", " \"\"\"Fast analysis of split quality using vectorized operations\"\"\"\n", " \n", " print(\"\\n\" + \"=\"*50)\n", " print(\"SPLIT QUALITY ANALYSIS\")\n", " print(\"=\"*50)\n", " \n", " # Duration - vectorized\n", " print(f\"\\nDuration: Train={train_df['duration_s'].mean():.1f}s, Test={test_df['duration_s'].mean():.1f}s\")\n", " \n", " # BPM - vectorized\n", " train_bpm = pd.to_numeric(train_df['bpm'], errors='coerce').mean()\n", " test_bpm = pd.to_numeric(test_df['bpm'], errors='coerce').mean()\n", " print(f\"BPM: Train={train_bpm:.1f}, Test={test_bpm:.1f}\")\n", " \n", " # Sample counts by key\n", " print(f\"\\nKey distribution:\")\n", " train_keys = train_df['key'].value_counts(normalize=True).head(5)\n", " test_keys = test_df['key'].value_counts(normalize=True).head(5)\n", " \n", " for key in train_keys.index[:3]:\n", " train_pct = train_keys.get(key, 0) * 100\n", " test_pct = test_keys.get(key, 0) * 100\n", " print(f\" {key}: Train={train_pct:.1f}%, Test={test_pct:.1f}%\")" ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": {}, "outputs": [], "source": [ "train_df, test_df = create_stratified_pack_split_fast(df_raw, n_processes=10, test_size=0.025)\n", "analyze_split_quality_fast(train_df, test_df)" ] }, { "cell_type": "code", "execution_count": null, "id": "5", "metadata": {}, "outputs": [], "source": [ "def create_combined_tags(row):\n", " bpm = row[\"bpm\"]\n", " key = row[\"key\"]\n", " processed_tags = row[\"processed_tags\"]\n", "\n", " tags = []\n", " if processed_tags is not None:\n", " tags = processed_tags.copy()\n", " if bpm is not None:\n", " tags.append(f\"{bpm} bpm\")\n", " if key is not None:\n", " tags.append([f\"key of {key}\"])\n", "\n", " return tags\n", "\n", "def process_df_for_encode(df):\n", " df_metas = df.copy()\n", " df_metas[\"id\"] = df_metas[\"uuid\"]\n", " df_metas['tags'] = df_metas.apply(create_combined_tags, axis=1)\n", " df_metas = df_metas.drop(columns=[\"uuid\", \"processed_tags\"])\n", "\n", " return df_metas" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "df_metas_val = process_df_for_encode(test_df)\n", "df_metas_val.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "7", "metadata": {}, "outputs": [], "source": [ "df_metas_tr = process_df_for_encode(train_df)\n", "df_metas_tr.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "df_metas_raw = process_df_for_encode(df_raw)\n", "df_metas_raw.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "9", "metadata": {}, "outputs": [], "source": [ "#upload_path = \"s3://suno-data/datasets/harvest/splice/sfx_metas_v0_all.jsonl\"\n", "#with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=True) as temp_file:\n", "# df_metas_raw.to_json(temp_file.name, orient='records', lines=True)\n", "# temp_filename = temp_file.name\n", "# results = upload_s3_files(from_local_filepaths=[temp_filename], to_s3_filepaths=[upload_path])\n", "#print(results)" ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "### Free Sound" ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "\n", "def extract_dominant_tag_batch(batch_data):\n", " \"\"\"Extract dominant tag for a batch of samples\"\"\"\n", " batch_id, df_batch = batch_data\n", " results = []\n", " \n", " pbar = tqdm(df_batch.iterrows(), total=len(df_batch), desc=f\"Batch {batch_id}\", leave=False)\n", " \n", " for idx, row in pbar:\n", " tags = row['tags_processed']\n", " \n", " # Simple handling: tags_processed is either a list of strings or None\n", " if tags is None or (isinstance(tags, list) and len(tags) == 0):\n", " dominant_tag = 'unknown'\n", " elif isinstance(tags, list):\n", " # Get most common tag from the list\n", " tag_counts = Counter(tags)\n", " dominant_tag = tag_counts.most_common(1)[0][0] if tag_counts else 'unknown'\n", " else:\n", " dominant_tag = 'unknown'\n", " \n", " results.append((idx, dominant_tag))\n", " \n", " pbar.close()\n", " return results\n", "\n", "def create_tag_prioritized_split_fast(df, test_size=0.2, random_state=42, n_processes=None):\n", " \"\"\"\n", " Fast stratified split prioritizing tag and duration balance over author grouping.\n", " \"\"\"\n", " if n_processes is None:\n", " n_processes = min(cpu_count(), 8)\n", " \n", " print(f\"Using {n_processes} processes\")\n", " print(f\"Total samples: {len(df)}\")\n", " \n", " # Extract dominant tag for each sample in parallel\n", " print(\"Extracting dominant tags for each sample...\")\n", " \n", " # Split dataframe into batches for parallel processing\n", " batch_size = max(1000, len(df) // (n_processes * 4)) # At least 1000 samples per batch\n", " df_batches = []\n", " for i in range(0, len(df), batch_size):\n", " batch = df.iloc[i:i + batch_size].copy()\n", " df_batches.append((len(df_batches), batch))\n", " \n", " # Process in parallel\n", " print(f\"Processing {len(df_batches)} batches across {n_processes} processes...\")\n", " \n", " with Pool(n_processes) as pool:\n", " batch_results = list(tqdm(\n", " pool.imap(extract_dominant_tag_batch, df_batches),\n", " total=len(df_batches),\n", " desc=\"Extracting tags\"\n", " ))\n", " \n", " # Flatten results and create mapping\n", " print(\"Assembling tag results...\")\n", " tag_mapping = {}\n", " for batch_result in batch_results:\n", " for idx, tag in batch_result:\n", " tag_mapping[idx] = tag\n", " \n", " # Add dominant tag to dataframe\n", " df_copy = df.copy()\n", " df_copy['dominant_tag'] = df_copy.index.map(tag_mapping)\n", " \n", " print(\"Creating stratification features...\")\n", " \n", " # Create duration bins\n", " df_copy['duration_bin'] = pd.cut(df_copy['duration_s'], bins=5, labels=['very_short', 'short', 'medium', 'long', 'very_long'])\n", " \n", " # Create stratification key\n", " df_copy['strat_key'] = df_copy['dominant_tag'].astype(str) + '_' + df_copy['duration_bin'].astype(str)\n", " \n", " # Analyze stratification viability\n", " print(\"Analyzing stratification options...\")\n", " strat_counts = df_copy['strat_key'].value_counts()\n", " tag_counts = df_copy['dominant_tag'].value_counts()\n", " duration_counts = df_copy['duration_bin'].value_counts()\n", " \n", " print(f\"Stratification analysis:\")\n", " print(f\" Unique tags: {df_copy['dominant_tag'].nunique()}\")\n", " print(f\" Unique tag+duration combinations: {len(strat_counts)}\")\n", " print(f\" Tag+duration strata with ≥10 samples: {(strat_counts >= 10).sum()}\")\n", " print(f\" Tag-only strata with ≥10 samples: {(tag_counts >= 10).sum()}\")\n", " \n", " # Choose stratification method\n", " try:\n", " # Try tag + duration stratification first\n", " viable_complex_strata = strat_counts[strat_counts >= 10]\n", " if len(viable_complex_strata) >= 10 and viable_complex_strata.sum() >= len(df) * 0.7:\n", " print(\"Using tag + duration stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['strat_key'],\n", " random_state=random_state\n", " )\n", " else:\n", " # Fall back to tag-only stratification\n", " viable_tag_strata = tag_counts[tag_counts >= 10]\n", " if len(viable_tag_strata) >= 5 and viable_tag_strata.sum() >= len(df) * 0.5:\n", " print(\"Using tag-only stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['dominant_tag'],\n", " random_state=random_state\n", " )\n", " else:\n", " # Fall back to duration-only stratification\n", " print(\"Using duration-only stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['duration_bin'],\n", " random_state=random_state\n", " )\n", " \n", " except Exception as e:\n", " print(f\"All stratification methods failed ({e}), using random split\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " random_state=random_state\n", " )\n", " \n", " # Remove helper columns\n", " columns_to_keep = [col for col in df.columns]\n", " train_df = train_df[columns_to_keep].copy()\n", " test_df = test_df[columns_to_keep].copy()\n", " \n", " print(f\"\\nSplit Results:\")\n", " print(f\"Train samples: {len(train_df)} ({len(train_df)/len(df)*100:.1f}%)\")\n", " print(f\"Test samples: {len(test_df)} ({len(test_df)/len(df)*100:.1f}%)\")\n", " \n", " # Check for author overlap if cap2 exists\n", " if 'cap2' in df.columns:\n", " train_authors = set(train_df['cap2'].dropna())\n", " test_authors = set(test_df['cap2'].dropna())\n", " overlap_authors = train_authors & test_authors\n", " if overlap_authors:\n", " print(f\"Warning: {len(overlap_authors)} authors appear in both train and test sets\")\n", " else:\n", " print(\"No author overlap between train and test sets\")\n", " \n", " return train_df, test_df\n", "\n", "def analyze_split_quality_fast(train_df, test_df):\n", " \"\"\"Fast analysis of split quality focusing on tag and duration balance\"\"\"\n", " \n", " print(\"\\n\" + \"=\"*50)\n", " print(\"SPLIT QUALITY ANALYSIS\")\n", " print(\"=\"*50)\n", " \n", " # Duration analysis\n", " print(f\"\\nDuration: Train={train_df['duration_s'].mean():.1f}s, Test={test_df['duration_s'].mean():.1f}s\")\n", " print(f\"Duration std: Train={train_df['duration_s'].std():.1f}s, Test={test_df['duration_s'].std():.1f}s\")\n", " \n", " # Author analysis if available\n", " if 'cap2' in train_df.columns:\n", " train_authors = train_df['cap2'].nunique()\n", " test_authors = test_df['cap2'].nunique()\n", " train_with_author = train_df['cap2'].notna().sum()\n", " test_with_author = test_df['cap2'].notna().sum()\n", " \n", " print(f\"\\nAuthors: Train={train_authors} unique, Test={test_authors} unique\")\n", " print(f\"Samples with authors: Train={train_with_author} ({train_with_author/len(train_df)*100:.1f}%), \"\n", " f\"Test={test_with_author} ({test_with_author/len(test_df)*100:.1f}%)\")\n", " \n", " # Top 10 tags analysis\n", " def get_top_tags(df, n=10):\n", " all_tags = []\n", " for tag_list in df['tags_processed'].dropna():\n", " if isinstance(tag_list, list):\n", " all_tags.extend(tag_list)\n", " return Counter(all_tags).most_common(n)\n", " \n", " print(f\"\\nTop 10 tags distribution:\")\n", " train_tags = dict(get_top_tags(train_df))\n", " test_tags = dict(get_top_tags(test_df))\n", " \n", " all_top_tags = set(train_tags.keys()) | set(test_tags.keys())\n", " for tag in list(all_top_tags)[:10]:\n", " train_pct = train_tags.get(tag, 0) / len(train_df) * 100\n", " test_pct = test_tags.get(tag, 0) / len(test_df) * 100\n", " diff = abs(train_pct - test_pct)\n", " print(f\" {tag:>15}: Train={train_pct:5.1f}%, Test={test_pct:5.1f}% (diff: {diff:.1f}%)\")\n", "\n", "# Example usage:\n", "# train_df, test_df = create_tag_prioritized_split_fast(df, n_processes=4)\n", "# analyze_split_quality_fast(train_df, test_df)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "train_df, test_df = create_tag_prioritized_split_fast(df_raw, n_processes=10, test_size=0.025)\n", "analyze_split_quality_fast(train_df, test_df)" ] }, { "cell_type": "code", "execution_count": null, "id": "13", "metadata": {}, "outputs": [], "source": [ "def create_combined_tags(row):\n", " cap1 = row[\"cap1\"]\n", " cap2 = row[\"cap2_filtered\"]\n", " processed_tags = row[\"tags_processed\"]\n", "\n", " tags = []\n", " if processed_tags is not None:\n", " tags = processed_tags.copy()\n", " if cap1 is not None:\n", " tags.append(cap1)\n", " if cap2 is not None:\n", " tags.append(cap2)\n", "\n", " return tags\n", "\n", "def process_df_for_encode(df):\n", " df_metas = df.copy()\n", " df_metas['tags'] = df_metas.apply(create_combined_tags, axis=1)\n", " df_metas = df_metas.drop(columns=[\"tags_processed\"])\n", "\n", " return df_metas" ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": {}, "outputs": [], "source": [ "df_metas_val = process_df_for_encode(test_df)\n", "df_metas_val.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": {}, "outputs": [], "source": [ "df_metas_tr = process_df_for_encode(train_df)\n", "df_metas_tr.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "16", "metadata": {}, "outputs": [], "source": [ "df_metas_raw = process_df_for_encode(df_raw)\n", "df_metas_raw.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "17", "metadata": {}, "outputs": [], "source": [ "upload_path = \"s3://suno-data/datasets/harvest/freesound/metas_v0_val.jsonl\"\n", "with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=True) as temp_file:\n", " df_metas_val.to_json(temp_file.name, orient='records', lines=True)\n", " temp_filename = temp_file.name\n", " results = upload_s3_files(from_local_filepaths=[temp_filename], to_s3_filepaths=[upload_path])\n", "print(results)" ] }, { "cell_type": "markdown", "id": "18", "metadata": {}, "source": [ "### Pond5 SFX" ] }, { "cell_type": "code", "execution_count": null, "id": "19", "metadata": {}, "outputs": [], "source": [ "def extract_dominant_tag_batch(batch_data):\n", " \"\"\"Extract dominant tag for a batch of samples\"\"\"\n", " batch_id, df_batch = batch_data\n", " results = []\n", " \n", " pbar = tqdm(df_batch.iterrows(), total=len(df_batch), desc=f\"Batch {batch_id}\", leave=False)\n", " \n", " for idx, row in pbar:\n", " tags = row['combined_tags']\n", " \n", " # Simple handling: combined_tags is either a list of strings or None\n", " if tags is None or (isinstance(tags, list) and len(tags) == 0):\n", " dominant_tag = 'unknown'\n", " elif isinstance(tags, list):\n", " # Get most common tag from the list\n", " tag_counts = Counter(tags)\n", " dominant_tag = tag_counts.most_common(1)[0][0] if tag_counts else 'unknown'\n", " else:\n", " dominant_tag = 'unknown'\n", " \n", " results.append((idx, dominant_tag))\n", " \n", " pbar.close()\n", " return results\n", "\n", "def create_tag_prioritized_split_fast(df, test_size=0.2, random_state=42, n_processes=None):\n", " \"\"\"\n", " Fast stratified split prioritizing tag, duration, and stereo balance over author grouping.\n", " \"\"\"\n", " if n_processes is None:\n", " n_processes = min(cpu_count(), 8)\n", " \n", " print(f\"Using {n_processes} processes\")\n", " print(f\"Total samples: {len(df)}\")\n", " print(f\"Total authors: {df['author'].nunique()}\")\n", " \n", " # Extract dominant tag for each sample in parallel\n", " print(\"Extracting dominant tags for each sample...\")\n", " \n", " # Split dataframe into batches for parallel processing\n", " batch_size = max(1000, len(df) // (n_processes * 4)) # At least 1000 samples per batch\n", " df_batches = []\n", " for i in range(0, len(df), batch_size):\n", " batch = df.iloc[i:i + batch_size].copy()\n", " df_batches.append((len(df_batches), batch))\n", " \n", " # Process in parallel\n", " print(f\"Processing {len(df_batches)} batches across {n_processes} processes...\")\n", " \n", " with Pool(n_processes) as pool:\n", " batch_results = list(tqdm(\n", " pool.imap(extract_dominant_tag_batch, df_batches),\n", " total=len(df_batches),\n", " desc=\"Extracting tags\"\n", " ))\n", " \n", " # Flatten results and create mapping\n", " print(\"Assembling tag results...\")\n", " tag_mapping = {}\n", " for batch_result in batch_results:\n", " for idx, tag in batch_result:\n", " tag_mapping[idx] = tag\n", " \n", " # Add dominant tag to dataframe\n", " df_copy = df.copy()\n", " df_copy['dominant_tag'] = df_copy.index.map(tag_mapping)\n", " \n", " print(\"Creating stratification features...\")\n", " \n", " # Create duration bins\n", " df_copy['duration_bin'] = pd.cut(df_copy['duration_s'], bins=4, \n", " labels=['very_short', 'short', 'medium', 'long'])\n", " \n", " # Create stereo category (handle None values)\n", " df_copy['stereo_category'] = df_copy['stereo'].fillna('unknown').astype(str)\n", " \n", " # Create stratification keys with different levels of complexity\n", " df_copy['strat_key_complex'] = (df_copy['dominant_tag'].astype(str) + '_' + \n", " df_copy['duration_bin'].astype(str) + '_' + \n", " df_copy['stereo_category'])\n", " \n", " df_copy['strat_key_medium'] = (df_copy['dominant_tag'].astype(str) + '_' + \n", " df_copy['duration_bin'].astype(str))\n", " \n", " df_copy['strat_key_simple'] = df_copy['dominant_tag'].astype(str)\n", " \n", " # Analyze stratification viability\n", " print(\"Analyzing stratification options...\")\n", " complex_counts = df_copy['strat_key_complex'].value_counts()\n", " medium_counts = df_copy['strat_key_medium'].value_counts()\n", " simple_counts = df_copy['strat_key_simple'].value_counts()\n", " stereo_counts = df_copy['stereo_category'].value_counts()\n", " duration_counts = df_copy['duration_bin'].value_counts()\n", " \n", " print(f\"Stratification analysis:\")\n", " print(f\" Tag+Duration+Stereo strata: {len(complex_counts)} total, {(complex_counts >= 10).sum()} viable (≥10 samples)\")\n", " print(f\" Tag+Duration strata: {len(medium_counts)} total, {(medium_counts >= 10).sum()} viable\")\n", " print(f\" Tag-only strata: {len(simple_counts)} total, {(simple_counts >= 10).sum()} viable\")\n", " print(f\" Stereo categories: {len(stereo_counts)} total, {(stereo_counts >= 10).sum()} viable\")\n", " print(f\" Duration bins: {len(duration_counts)} total, {(duration_counts >= 10).sum()} viable\")\n", " \n", " # Choose stratification method\n", " try:\n", " # Try tag + duration + stereo stratification first\n", " viable_complex_strata = complex_counts[complex_counts >= 10]\n", " if len(viable_complex_strata) >= 10 and viable_complex_strata.sum() >= len(df) * 0.7:\n", " print(\"Using tag + duration + stereo stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['strat_key_complex'],\n", " random_state=random_state\n", " )\n", " else:\n", " # Fall back to tag + duration stratification\n", " viable_medium_strata = medium_counts[medium_counts >= 10]\n", " if len(viable_medium_strata) >= 8 and viable_medium_strata.sum() >= len(df) * 0.6:\n", " print(\"Using tag + duration stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['strat_key_medium'],\n", " random_state=random_state\n", " )\n", " else:\n", " # Fall back to tag-only stratification\n", " viable_simple_strata = simple_counts[simple_counts >= 10]\n", " if len(viable_simple_strata) >= 5 and viable_simple_strata.sum() >= len(df) * 0.5:\n", " print(\"Using tag-only stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['strat_key_simple'],\n", " random_state=random_state\n", " )\n", " else:\n", " # Fall back to duration-only stratification\n", " print(\"Using duration-only stratification\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " stratify=df_copy['duration_bin'],\n", " random_state=random_state\n", " )\n", " \n", " except Exception as e:\n", " print(f\"All stratification methods failed ({e}), using random split\")\n", " train_df, test_df = train_test_split(\n", " df_copy,\n", " test_size=test_size,\n", " random_state=random_state\n", " )\n", " \n", " # Remove helper columns\n", " columns_to_keep = [col for col in df.columns]\n", " train_df = train_df[columns_to_keep].copy()\n", " test_df = test_df[columns_to_keep].copy()\n", " \n", " print(f\"\\nSplit Results:\")\n", " print(f\"Train samples: {len(train_df)} ({len(train_df)/len(df)*100:.1f}%)\")\n", " print(f\"Test samples: {len(test_df)} ({len(test_df)/len(df)*100:.1f}%)\")\n", " \n", " # Check for author overlap\n", " train_authors = set(train_df['author'].dropna())\n", " test_authors = set(test_df['author'].dropna())\n", " overlap_authors = train_authors & test_authors\n", " if overlap_authors:\n", " print(f\"Warning: {len(overlap_authors)} authors appear in both train and test sets\")\n", " print(f\" Total unique authors: train={len(train_authors)}, test={len(test_authors)}\")\n", " else:\n", " print(\"No author overlap between train and test sets\")\n", " \n", " return train_df, test_df\n", "\n", "def analyze_split_quality_fast(train_df, test_df):\n", " \"\"\"Fast analysis of split quality focusing on tag, duration, and stereo balance\"\"\"\n", " \n", " print(\"\\n\" + \"=\"*60)\n", " print(\"SPLIT QUALITY ANALYSIS\")\n", " print(\"=\"*60)\n", " \n", " # Duration analysis\n", " print(f\"\\nDURATION DISTRIBUTION:\")\n", " train_dur = train_df['duration_s']\n", " test_dur = test_df['duration_s']\n", " \n", " print(f\" Mean: Train={train_dur.mean():.1f}s, Test={test_dur.mean():.1f}s\")\n", " print(f\" Std: Train={train_dur.std():.1f}s, Test={test_dur.std():.1f}s\")\n", " print(f\" Median: Train={train_dur.median():.1f}s, Test={test_dur.median():.1f}s\")\n", " \n", " # Stereo analysis\n", " print(f\"\\nSTEREO DISTRIBUTION:\")\n", " train_stereo = train_df['stereo'].value_counts(normalize=True, dropna=False)\n", " test_stereo = test_df['stereo'].value_counts(normalize=True, dropna=False)\n", " \n", " all_stereo_values = set(train_stereo.index) | set(test_stereo.index)\n", " for stereo_val in sorted(all_stereo_values, key=str):\n", " train_pct = train_stereo.get(stereo_val, 0) * 100\n", " test_pct = test_stereo.get(stereo_val, 0) * 100\n", " diff = abs(train_pct - test_pct)\n", " print(f\" {str(stereo_val):>8}: Train={train_pct:5.1f}%, Test={test_pct:5.1f}% (diff: {diff:.1f}%)\")\n", " \n", " # Author analysis\n", " print(f\"\\nAUTHOR DISTRIBUTION:\")\n", " train_authors = train_df['author'].nunique()\n", " test_authors = test_df['author'].nunique()\n", " total_authors = pd.concat([train_df['author'], test_df['author']]).nunique()\n", " \n", " print(f\" Unique authors: Train={train_authors}, Test={test_authors}, Total={total_authors}\")\n", " \n", " # Author overlap analysis\n", " train_author_set = set(train_df['author'].dropna())\n", " test_author_set = set(test_df['author'].dropna())\n", " overlap = len(train_author_set & test_author_set)\n", " if overlap > 0:\n", " print(f\" Author overlap: {overlap} authors appear in both sets\")\n", " \n", " # Top 10 tags analysis\n", " def get_top_tags(df, n=10):\n", " all_tags = []\n", " for tag_list in df['combined_tags'].dropna():\n", " if isinstance(tag_list, list):\n", " all_tags.extend(tag_list)\n", " return Counter(all_tags).most_common(n)\n", " \n", " print(f\"\\nTOP 10 TAGS DISTRIBUTION:\")\n", " train_tags = dict(get_top_tags(train_df))\n", " test_tags = dict(get_top_tags(test_df))\n", " \n", " all_top_tags = set(train_tags.keys()) | set(test_tags.keys())\n", " for tag in list(all_top_tags)[:10]:\n", " train_pct = train_tags.get(tag, 0) / len(train_df) * 100\n", " test_pct = test_tags.get(tag, 0) / len(test_df) * 100\n", " diff = abs(train_pct - test_pct)\n", " print(f\" {tag:>15}: Train={train_pct:5.1f}%, Test={test_pct:5.1f}% (diff: {diff:.1f}%)\")" ] }, { "cell_type": "code", "execution_count": null, "id": "20", "metadata": {}, "outputs": [], "source": [ "train_df, test_df = create_tag_prioritized_split_fast(df_raw, n_processes=10, test_size=0.025)\n", "analyze_split_quality_fast(train_df, test_df)" ] }, { "cell_type": "code", "execution_count": null, "id": "21", "metadata": {}, "outputs": [], "source": [ "def create_combined_tags(row):\n", " genre = row[\"genre\"]\n", " description = row[\"description\"]\n", " name = row[\"name\"]\n", " taxonomy = row[\"taxonomy\"]\n", " stereo = row[\"stereo\"]\n", " combined_tags = row[\"combined_tags\"]\n", "\n", " tags = []\n", " if combined_tags is not None:\n", " tags = combined_tags.copy()\n", " if genre is not None:\n", " tags.append(genre)\n", " if description is not None:\n", " tags.append(description)\n", " if name is not None:\n", " tags.append(name)\n", " if taxonomy is not None:\n", " for tax in taxonomy:\n", " tags.append(tax)\n", " if stereo is not None:\n", " if stereo:\n", " tags.append(\"stereo\")\n", " else:\n", " tags.append(\"mono\")\n", "\n", " return tags\n", "\n", "def process_df_for_encode(df):\n", " df_metas = df.copy()\n", " df_metas['tags'] = df_metas.apply(create_combined_tags, axis=1)\n", " df_metas = df_metas.drop(columns=[\"combined_tags\"])\n", "\n", " return df_metas" ] }, { "cell_type": "code", "execution_count": null, "id": "22", "metadata": {}, "outputs": [], "source": [ "df_metas_val = process_df_for_encode(test_df)\n", "df_metas_val.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "23", "metadata": {}, "outputs": [], "source": [ "df_metas_tr = process_df_for_encode(train_df)\n", "df_metas_tr.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "24", "metadata": {}, "outputs": [], "source": [ "df_metas_raw = process_df_for_encode(df_raw)\n", "df_metas_raw.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "25", "metadata": {}, "outputs": [], "source": [ "upload_path = \"s3://suno-data/datasets/harvest/pond5_sfx/pond5_metas_v0_val.jsonl\"\n", "with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=True) as temp_file:\n", " df_metas_val.to_json(temp_file.name, orient='records', lines=True)\n", " temp_filename = temp_file.name\n", " results = upload_s3_files(from_local_filepaths=[temp_filename], to_s3_filepaths=[upload_path])\n", "print(results)" ] }, { "cell_type": "markdown", "id": "26", "metadata": {}, "source": [ "### Make Silence for Padding" ] }, { "cell_type": "code", "execution_count": null, "id": "27", "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.dac_vae_fixed_25hz import preload_models\n", "from suno_utils.tasks.dac_vae_fixed_25hz import encode_overlap as encode, SAMPLE_RATE\n", "import torch\n", "\n", "model_filepath = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n", "\n", "_ = preload_models(\n", " checkpoint_filepath=model_filepath,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "28", "metadata": {}, "outputs": [], "source": [ "silence_15sec = torch.zeros((2, 15 * SAMPLE_RATE), dtype=torch.float32)" ] }, { "cell_type": "code", "execution_count": null, "id": "29", "metadata": {}, "outputs": [], "source": [ "encoded_arrays = []\n", "for nn, arr in enumerate([silence_15sec]):\n", " encoded_arrays.append(encode([arr], normalize_volume=False)[0])" ] }, { "cell_type": "code", "execution_count": null, "id": "30", "metadata": {}, "outputs": [], "source": [ "np.savez(\"/home/sara/dac_vae_tuned_25hz_15s_silence.npz\", **{\"vae_data\": encoded_arrays[0]})" ] }, { "cell_type": "code", "execution_count": null, "id": "31", "metadata": {}, "outputs": [], "source": [ "silence_vae = np.load(\"/home/sara/dac_vae_tuned_25hz_15s_silence.npz\")[\"vae_data\"]" ] }, { "cell_type": "code", "execution_count": null, "id": "32", "metadata": {}, "outputs": [], "source": [ "silence_vae.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "33", "metadata": {}, "outputs": [], "source": [ "silence_vae" ] }, { "cell_type": "markdown", "id": "34", "metadata": {}, "source": [ "## More Filtering" ] }, { "cell_type": "code", "execution_count": null, "id": "35", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import read_jsonl\n", "import random\n", "import pandas as pd\n", "import re" ] }, { "cell_type": "code", "execution_count": null, "id": "36", "metadata": {}, "outputs": [], "source": [ "def summarize_meta(ds_path):\n", " metas = read_jsonl(ds_path)\n", " print(f\"{len(metas):,} tracks with {sum([m['duration_s'] for m in metas])/60/60:,.1f}h total\")\n", " print(metas[0].keys())\n", " return metas\n", "\n", "def sample_meta(metas):\n", " random_index = random.randint(0, len(metas ) - 1)\n", " sample = metas[random_index]\n", " for k,v in sample.items():\n", " print(f\"{k}: {v}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "37", "metadata": {}, "outputs": [], "source": [ "metas_tr = summarize_meta(\"/app2/suno/data/diffusion/sfx/v0/combined_metas.jsonl\")" ] }, { "cell_type": "code", "execution_count": null, "id": "38", "metadata": {}, "outputs": [], "source": [ "sample_meta(metas_tr)" ] }, { "cell_type": "code", "execution_count": null, "id": "39", "metadata": {}, "outputs": [], "source": [ "df = pd.DataFrame.from_dict(metas_tr)\n", "df.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "40", "metadata": {}, "outputs": [], "source": [ "url_pattern = r'(https?://[^\\s\\]\\)]+|www\\.[^\\s\\]\\)]+)'\n", "\n", "def filter_tags(tags):\n", " if not isinstance(tags, list) or len(tags) == 0:\n", " return tags # this shouldn't happen\n", " \n", " filtered_tags = []\n", "\n", " for tag in tags:\n", " urls = re.findall(url_pattern, tag)\n", " if len(urls) > 0:\n", " continue\n", " else:\n", " filtered_tags.append(tag)\n", "\n", " return filtered_tags\n", "\n", "\n", "df['tags'] = df['tags'].apply(filter_tags)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "41", "metadata": {}, "outputs": [], "source": [ "def filter_tags(tags):\n", " if not isinstance(tags, list) or len(tags) == 0:\n", " return tags # this shouldn't happen\n", "\n", " filtered_tags = []\n", "\n", " for tag in tags:\n", " if \"https\" in tag or \"http\" in tag or \"www.\" in tag or \".com\" in tag:\n", " continue\n", " else:\n", " filtered_tags.append(tag)\n", "\n", " return filtered_tags\n", " \n", "df['tags'] = df['tags'].apply(filter_tags)" ] }, { "cell_type": "code", "execution_count": null, "id": "42", "metadata": {}, "outputs": [], "source": [ "email_pattern = r'\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}\\b'\n", "\n", "def filter_tags(tags):\n", " if not isinstance(tags, list) or len(tags) == 0:\n", " return tags # this shouldn't happen\n", "\n", " filtered_tags = []\n", "\n", " for tag in tags:\n", " sentences = re.split(r'(?<=[.!?])\\s+', tag.strip())\n", " filtered_sentences = [s for s in sentences if not re.search(email_pattern, s)]\n", " if len(filtered_sentences) > 0:\n", " filtered_text = ' '.join(filtered_sentences)\n", " filtered_tags.append(filtered_text)\n", "\n", " return filtered_tags\n", "df['tags'] = df['tags'].apply(filter_tags)" ] }, { "cell_type": "code", "execution_count": null, "id": "43", "metadata": {}, "outputs": [], "source": [ "def filter_tags(tags):\n", " if not isinstance(tags, list) or len(tags) == 0:\n", " return tags # this shouldn't happen\n", "\n", " filtered_tags = []\n", "\n", " def do_filter(tag):\n", " if \"please\" in tag.lower() or \"contact me\" in tag.lower():\n", " return True\n", " else:\n", " return False\n", "\n", " for tag in tags:\n", " sentences = re.split(r'(?<=[.!?])\\s+', tag.strip())\n", " filtered_sentences = [s for s in sentences if not do_filter(s)]\n", " if len(filtered_sentences) > 0:\n", " filtered_text = ' '.join(filtered_sentences)\n", " filtered_tags.append(filtered_text)\n", "\n", " return filtered_tags\n", "\n", "df['tags'] = df['tags'].apply(filter_tags)" ] }, { "cell_type": "code", "execution_count": null, "id": "44", "metadata": {}, "outputs": [], "source": [ "for i in range(10):\n", " sample = df.sample(n=1).iloc[0]\n", " print(sample.tags)" ] }, { "cell_type": "code", "execution_count": null, "id": "45", "metadata": {}, "outputs": [], "source": [ "df.head()" ] }, { "cell_type": "code", "execution_count": null, "id": "46", "metadata": {}, "outputs": [], "source": [ "df.to_json(\"combined_metas.jsonl\", orient='records', lines=True)" ] }, { "cell_type": "code", "execution_count": null, "id": "47", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_clean", "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.10.15" } }, "nbformat": 4, "nbformat_minor": 5 }