{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from suno_utils.utils.text import read_json, read_jsonl\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "old_ids_path = \"/app2/suno/data/auk_v0/ids_keep_sets_ext_v11.json\"\n",
    "\n",
    "old_ids = json.load(open(old_ids_path))\n",
    "for key, value in old_ids.items():\n",
    "    print(f\"{key}: {len(value):,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "metas_path = \"/app2/suno/data/auk_v0/metas_v4_tr.jsonl\"\n",
    "metadata = read_jsonl(metas_path)\n",
    "print(metadata[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_ids = old_ids[\"audio_ids\"]\n",
    "print(generation_ids[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadata_dict = {meta[\"id\"]: meta for meta in metadata}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "sft_rows = []\n",
    "for gen_id in generation_ids:\n",
    "    if gen_id in metadata_dict:\n",
    "        sft_rows.append(metadata_dict[gen_id])\n",
    "    else:\n",
    "        print(\"skipping\", gen_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(sft_rows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(sft_rows)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv('sft_data.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('sft_data.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter = df[df['duration_s'] > 180]\n",
    "df_filter = df_filter[df_filter['lang'] == 'en']\n",
    "df_filter = df_filter[df_filter['tags'].notna()]\n",
    "df_filter = df_filter[df_filter['tags'].apply(len) > 100]\n",
    "df_filter = df_filter[df_filter['weight'] > 1.2]\n",
    "len(df_filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "genres = ['jazz', 'hip hop', 'country', 'folk', 'electronic', 'pop', 'classical', 'rock', 'metal']\n",
    "\n",
    "def find_genres_in_tags(tags_list, genres_list):\n",
    "    \"\"\"\n",
    "    Find which genres appear in the tags list.\n",
    "    Returns a list of genres that appear at least once.\n",
    "    \"\"\"\n",
    "    print(type(tags_list))\n",
    "    if pd.isna(tags_list):\n",
    "        return []\n",
    "    \n",
    "    found_genres = []\n",
    "    # Convert tags to lowercase for case-insensitive matching\n",
    "    tags_lower = [tag.lower() for tag in tags_list]\n",
    "    tags_string = ' '.join(tags_lower)  # Join all tags into one string for easier searching\n",
    "    print(tags_string)\n",
    "    \n",
    "    for genre in genres_list:\n",
    "        # Check if genre appears as a substring in any of the tags\n",
    "        if genre.lower() in tags_string:\n",
    "            found_genres.append(genre)\n",
    "    \n",
    "    return found_genres\n",
    "\n",
    "# Apply the function to create a new column\n",
    "df_filter['found_genres'] = df_filter['tags'].apply(lambda x: find_genres_in_tags(x, genres))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_filter = df_filter[df_filter['found_genres'].apply(len) > 0]\n",
    "len(df_filter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "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
}