{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import polars as pl\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we load the original metas to get the artists and song titles\n",
    "METAS_DIR = \"/app/suno/tmp\"\n",
    "#raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_discogs_subset_metas.jsonl\"))\n",
    "raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_genius_metas.jsonl\"))\n",
    "#raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_imslp_metas.jsonl\"))\n",
    "print(len(raw_metas))\n",
    "#dataset_name = \"imslp\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"genius\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas_filtered = [m for m in raw_metas if m[\"views\"] > 100_000]\n",
    "print(len(raw_metas_filtered))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_metas_filtered[7]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " # filter the raw metas based on the ids in our finetuning set\n",
    "import json\n",
    "from tqdm import tqdm\n",
    "#filepath = \"/home/christian/code/christian/metadata/v45_splits/info_tr_ft_v10_ids.json\"\n",
    "filepath = \"/home/christian/code/christian/metadata/v45_splits/ids_keep_sets_v11.json\"\n",
    "with open(filepath, \"r\") as f:\n",
    "    info = json.load(f)\n",
    "\n",
    "print(info.keys())\n",
    "\n",
    "\n",
    "# get the ids from the info dict\n",
    "ids_keep = set(info[dataset_name])\n",
    "print(len(ids_keep))\n",
    "\n",
    "\n",
    "filtered_metas = []\n",
    "for meta in tqdm(raw_metas):\n",
    "    #if meta[\"id\"] in ids_keep:\n",
    "    filtered_metas.append(meta)\n",
    "\n",
    "print(len(filtered_metas))\n",
    "\n",
    "if dataset_name == \"discogs_subset\":\n",
    "    # lets filter harder, remove any that don't have a discogs release date\n",
    "    #filtered_metas = [meta for meta in filtered_metas if meta.get(\"discogs\", {}).get(\"released\", \"unknown\") != \"unknown\"]\n",
    "    #print(len(filtered_metas))\n",
    "\n",
    "    # lets filter harder, remove any that don't have a rym genre\n",
    "    #filtered_metas = [meta for meta in filtered_metas if meta.get(\"rym_genres\", [])]\n",
    "    #print(len(filtered_metas))\n",
    "\n",
    "    # also filter on views\n",
    "    filtered_metas = [meta for meta in filtered_metas if meta.get(\"views\", 0) > 125_000]\n",
    "    print(len(filtered_metas))\n",
    "\n",
    "elif dataset_name == \"imslp\":\n",
    "    # filter if no title\n",
    "    filtered_metas = [meta for meta in filtered_metas if \"title\" in meta]\n",
    "    print(len(filtered_metas))\n",
    "\n",
    "elif dataset_name == \"genius\":\n",
    "    # filter if no title\n",
    "    filtered_metas = [meta for meta in filtered_metas if \"-\" in meta[\"youtube_title\"]]\n",
    "    filtered_metas = [meta for meta in filtered_metas if \" - \" in meta[\"youtube_title\"]]\n",
    "    print(len(filtered_metas))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = filtered_metas[10]\n",
    "\n",
    "song_info = f\"\"\"\n",
    "title: {meta[\"title\"]}\n",
    "composer: {meta[\"composer\"]}\n",
    "existing tags: {meta.get(\"tags\", [])}\n",
    "\"\"\".strip()\n",
    "\n",
    "print(song_info)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta = filtered_metas[5022]\n",
    "\n",
    "title = meta[\"title\"]\n",
    "artist = meta[\"artists\"][0][\"name\"]\n",
    "rym_genres = meta.get(\"rym_genres\", [])\n",
    "discogs = meta.get(\"discogs\", {})\n",
    "year = discogs.get(\"released\", \"unknown\")\n",
    "\n",
    "print(f\"title: {title}\")\n",
    "print(f\"artist: {artist}\")\n",
    "print(f\"year: {year}\")\n",
    "print(f\"existing tags: {rym_genres}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "for n in range(10):\n",
    "    meta = filtered_metas[np.random.randint(0, len(filtered_metas))]\n",
    "    song_info = f\"\"\"{n}: {meta[\"youtube_title\"]}\"\"\"\n",
    "    #existing tags: {meta.get(\"genius_tags\", [])[:10]}\n",
    "    #\"\"\"\n",
    "    print(song_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "client = OpenAI(api_key=\"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\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a music expert with an expansive knowledge of music.\n",
    "You will be provided with a song, the artist, and some descriptors of the song.\n",
    "Try to describe the main elements of the song both musically and acoustically, and note any unique qualities, sound effects, or elements. \n",
    "Make sure to include all the instruments present. \n",
    "Do not include the name of any artists or musicians or refer to the name of any songs in your response. \n",
    "Do not respond in complete sentences, just to provide a caption-like response as a single block of text. \n",
    "If this song is a duet, please note that in your response.\n",
    "If this song has charted or won awards, please note that in your response.\n",
    "If the song is a classical piece respond with a description of the evolution of the piece overtime and the instruments.\n",
    "If you are not familiar with the song respond with \"None\":\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are an expert audio engineer.\n",
    "Given a list of songs (title and artist), provide a score from 0 to 100 that reflects the audio production quality of each recording.\n",
    "\n",
    "Use the entire 0–100 scale.\n",
    "Penalize recordings with:\n",
    "Distortion\n",
    "Poor mixing or imbalances\n",
    "Over-compression or lack of dynamic range\n",
    "Muddy or harsh tonal quality\n",
    "\n",
    "Favor recordings with:\n",
    "Clarity and balance\n",
    "Dynamic range\n",
    "Professional mixing/mastering\n",
    "High-fidelity sonic detail\n",
    "\n",
    "Return your response in this exact JSON format:\n",
    "{\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Steely Dan - Aja\",\n",
    "    \"score\": 100,\n",
    "    \"reasoning\": \"Audiophile benchmark; immaculate recording and mix\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Dire Straits - Brothers in Arms\",\n",
    "    \"score\": 98,\n",
    "    \"reasoning\": \"Crystal-clear digital recording, great stereo imaging\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Daft Punk - Giorgio by Moroder\",\n",
    "    \"score\": 95,\n",
    "    \"reasoning\": \"Clean layering and dynamic range in complex mix\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Norah Jones - Don't Know Why\",\n",
    "    \"score\": 94,\n",
    "    \"reasoning\": \"Warm, intimate acoustic production\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Metallica - ...And Justice for All\",\n",
    "    \"score\": 20,\n",
    "    \"reasoning\": \"Iconic failure in mixing; total lack of bass guitar\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Red Hot Chili Peppers - Californication\",\n",
    "    \"score\": 25,\n",
    "    \"reasoning\": \"Brutal loudness war victim; clipping and distortion everywhere\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"The Strokes - Is This It (UK version)\",\n",
    "    \"score\": 30,\n",
    "    \"reasoning\": \"Intentionally lo-fi, but still grating and tinny\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Oasis - Be Here Now\",\n",
    "    \"score\": 22,\n",
    "    \"reasoning\": \"Wall of sludge; no headroom, bloated beyond listenability\"\n",
    "  },\n",
    "  \"<id>\": {\n",
    "    \"song\": \"Green Day - American Idiot\",\n",
    "    \"score\": 35,\n",
    "    \"reasoning\": \"Radio-ready but painfully flat; hypercompression kills dynamics\"\n",
    "  }\n",
    "}\n",
    "If you are unfamiliar with a song, respond with:\n",
    "\n",
    "\"score\": null\n",
    "\"reasoning\" null\n",
    "\n",
    "Do not include any commentary outside the JSON.\n",
    "Do not refer to artist names or song titles in your reasoning.\n",
    "Your entire response must be in JSON using the format above.\n",
    "\n",
    "The songs to score are:\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "\n",
    "You are an expert in music understanding and natural language generation.\n",
    "Your task is to take a short caption describing a piece of music and expand it into multiple persona-based descriptions.\n",
    "\n",
    "You will be given:\n",
    "- An original caption describing a piece of music (short text).\n",
    "- Optionally, metadata including the song title and artist name.\n",
    "\n",
    "Use this metadata as contextual inspiration to better infer mood, style, and musical details — but you must never include or reference the artist name, album name, or song title in the output. \n",
    "Treat them only as hidden hints for tone and style.\n",
    "\n",
    "Your goal:\n",
    "Generate rewritten versions of the original caption for each of five personas, in three lengths:\n",
    "- short (~1 sentence)\n",
    "- medium (~2–3 sentences)\n",
    "- long (~up to 1000 characters)\n",
    "\n",
    "Each rewrite should:\n",
    "- Reflect the persona’s distinct voice, focus, and vocabulary style.\n",
    "- Remain faithful to the content of the original caption.\n",
    "- Avoid hallucinating new musical elements that contradict the original description.\n",
    "- Never mention or allude to any real names, albums, or titles.\n",
    "\n",
    "---\n",
    "\n",
    "### Personas\n",
    "\n",
    "1. **Casual Listener**\n",
    "   - Focus: Mood, vibe, emotional feel.\n",
    "   - Style: Intuitive, plain, and context-driven; uses everyday emotional or situational imagery.\n",
    "\n",
    "2. **Producer**\n",
    "   - Focus: Instrumentation, rhythm, and mix.\n",
    "   - Style: Technical but concise; uses production, arrangement, and sound-engineering vocabulary.\n",
    "\n",
    "3. **Composer**\n",
    "   - Focus: Harmony, form, orchestration, and structure.\n",
    "   - Style: Analytical and musical; references keys, progressions, and formal organization.\n",
    "\n",
    "4. **Film Scorer**\n",
    "   - Focus: Narrative arc, emotion, and cinematic context.\n",
    "   - Style: Descriptive, scene-based; links the sound to an emotional or visual story.\n",
    "\n",
    "5. **Sound Designer**\n",
    "   - Focus: Texture, timbre, and motion of sound.\n",
    "   - Style: Abstract, sensory, and process-oriented; uses tactile or spatial metaphors.\n",
    "\n",
    "---\n",
    "\n",
    "### Input format:\n",
    "original caption: [your input caption here]  \n",
    "artist: [optional artist name]  \n",
    "title: [optional song title]\n",
    "\n",
    "---\n",
    "\n",
    "### Output format:\n",
    "Return a single valid JSON object with the following structure:\n",
    "\n",
    "{\n",
    "  \"original_caption\": \"<input caption>\",\n",
    "  \"artist\": \"<input artist or null>\",\n",
    "  \"title\": \"<input title or null>\",\n",
    "  \"augmented_captions\": {\n",
    "    \"Casual Listener\": { \"short\": \"...\", \"medium\": \"...\", \"long\": \"...\" },\n",
    "    \"Producer\": { \"short\": \"...\", \"medium\": \"...\", \"long\": \"...\" },\n",
    "    \"Composer\": { \"short\": \"...\", \"medium\": \"...\", \"long\": \"...\" },\n",
    "    \"Film Scorer\": { \"short\": \"...\", \"medium\": \"...\", \"long\": \"...\" },\n",
    "    \"Sound Designer\": { \"short\": \"...\", \"medium\": \"...\", \"long\": \"...\" }\n",
    "  }\n",
    "}\n",
    "\n",
    "Rules:\n",
    "- Output valid JSON only — no markdown, commentary, or preamble.\n",
    "- Do not reference or imply the artist, title, or album in any form.\n",
    "- The “long” descriptions must not exceed 1000 characters.\n",
    "- Stay semantically faithful to the original caption.\n",
    "- When metadata is provided, you may use it as stylistic guidance (e.g., infer likely genre or tone) but never reveal it.\n",
    "\n",
    "---\n",
    "\n",
    "### Example\n",
    "\n",
    "original caption: [dreamy electronic track with soft vocals]  \n",
    "artist: [Beach House]  \n",
    "title: [Space Song]\n",
    "\n",
    "Expected output:\n",
    "\n",
    "{\n",
    "  \"original_caption\": \"dreamy electronic track with soft vocals\",\n",
    "  \"artist\": \"Beach House\",\n",
    "  \"title\": \"Space Song\",\n",
    "  \"augmented_captions\": {\n",
    "    \"Casual Listener\": {\n",
    "      \"short\": \"A dreamy, relaxing song that feels warm and nostalgic.\",\n",
    "      \"medium\": \"A calm and soothing electronic song with soft vocals, perfect for late-night listening or rainy evenings when everything feels still.\",\n",
    "      \"long\": \"This track feels like floating through a quiet city at night while neon lights blur in the rain. Soft synths and gentle vocals create a sense of calm and longing, blending electronic textures with a warm emotional glow. It’s peaceful yet full of unspoken feeling, like remembering something beautiful that’s just out of reach.\"\n",
    "    },\n",
    "    \"Producer\": {\n",
    "      \"short\": \"Smooth electronic track with soft synth pads and airy vocals.\",\n",
    "      \"medium\": \"Warm synth layers and a slow, sidechained beat support airy vocals processed with subtle reverb and stereo shimmer for a lush mix.\",\n",
    "      \"long\": \"Built around analog-style pads and a gentle pulsing kick, the production uses wide reverb tails and delicate compression to keep the sound light and open. The vocal sits mid-forward with a soft high-shelf lift, adding clarity without harshness. Everything is balanced and restrained, designed for depth and warmth rather than intensity.\"\n",
    "    },\n",
    "    \"Composer\": {\n",
    "      \"short\": \"In A minor with a circular four-chord progression and airy synth textures.\",\n",
    "      \"medium\": \"A slow electronic piece in A minor, cycling through a gentle four-chord loop beneath sustained synth harmonies and a simple melodic motif.\",\n",
    "      \"long\": \"Structured in loose ternary form, the piece centers on a repeating four-chord progression in A minor that maintains suspended tension. Layered synth voices act as a harmonic pad, while a floating vocal melody traces long arcs above the texture. Sparse rhythmic motion emphasizes harmonic stasis, creating a meditative, nocturnal sound world.\"\n",
    "    },\n",
    "    \"Film Scorer\": {\n",
    "      \"short\": \"Music for a quiet dawn scene filled with calm and hope.\",\n",
    "      \"medium\": \"Gentle synth tones rise like the first light of morning, paired with soft vocals that suggest quiet optimism and renewal.\",\n",
    "      \"long\": \"Imagine a scene where two people reunite after a long night apart. The music begins as soft pulses under a pale sky, slowly brightening with layers of harmony. Vocals enter like whispered dialogue, expressing emotion without words. Each phrase builds and fades like breath, underscoring a moment of serenity and fragile hope.\"\n",
    "    },\n",
    "    \"Sound Designer\": {\n",
    "      \"short\": \"Soft clouds of synth texture drift in slow motion.\",\n",
    "      \"medium\": \"A swirl of airy pads and dissolving vocals that melt into shimmering overtones, like fog illuminated by distant lights.\",\n",
    "      \"long\": \"The piece unfolds as a shifting atmosphere — glassy tones blur and merge while faint harmonics flicker at the edges. The voice becomes texture, elongated and spectral, weaving into vaporous synth layers. Space feels liquid, boundaries soft, as if the air itself were resonating. It’s more sensation than song — a living sonic mist.\"\n",
    "    }\n",
    "  }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "\n",
    "MAX_METAS = 400_000\n",
    "# Define the results file path\n",
    "results_file = f'/home/christian/code/christian/metadata/tagging/gpt-4_1-audio-tagging_results_{dataset_name}.json'\n",
    "\n",
    "# Initialize results dictionary\n",
    "results = {}\n",
    "\n",
    "# Load existing results if file exists\n",
    "if os.path.exists(results_file):\n",
    "    with open(results_file, 'r') as f:\n",
    "        results = json.load(f)\n",
    "    print(f\"Loaded {len(results)} existing results from {results_file}\")\n",
    "\n",
    "# Filter out metas that have already been processed\n",
    "metas_to_process = [meta for meta in filtered_metas[:MAX_METAS] if meta[\"id\"] not in results]\n",
    "print(f\"Processing {len(metas_to_process)} new items out of {len(filtered_metas[:MAX_METAS])} total\")\n",
    "\n",
    "\n",
    "# we will create a jsonl file for batch api\n",
    "batch_file = []\n",
    "\n",
    "# split the metas in to list of chunk_size\n",
    "chunk_size = 50\n",
    "chunked_metas = [metas_to_process[i:i+chunk_size] for i in range(0, len(metas_to_process), chunk_size)]\n",
    "print(len(chunked_metas))\n",
    "\n",
    "# to start lets just do 10 chunks\n",
    "\n",
    "for chunk in chunked_metas[:10]:\n",
    "    chunk_info = \"\"\n",
    "    for idx, meta in enumerate(chunk):\n",
    "        if dataset_name == \"discogs_subset\":\n",
    "            song_info = f\"\"\"\n",
    "            title: {meta[\"title\"]}\n",
    "            artist: {meta[\"artists\"][0][\"name\"]}\n",
    "            year: {meta.get(\"discogs\", {}).get(\"released\", \"unknown\")}\n",
    "            existing tags: {meta.get(\"rym_genres\", [])}\n",
    "            \"\"\".strip()\n",
    "        elif dataset_name == \"genius\":\n",
    "            #song_info = f\"\"\"\n",
    "            #title: {meta[\"youtube_title\"]} ({meta.get(\"genius_slug\", \"\")})\n",
    "            #existing tags: {meta.get(\"genius_tags\", [])[:10]}\n",
    "            #\"\"\".strip()\n",
    "            song_info = f\"\"\"{meta[\"id\"]}: {meta[\"youtube_title\"]}\"\"\".strip()\n",
    "        elif dataset_name == \"imslp\":\n",
    "            song_info = f\"\"\"\n",
    "            title: {meta[\"title\"]}\n",
    "            composer: {meta[\"composer\"]}\n",
    "            existing tags: {meta.get(\"tags\", [])}\n",
    "            \"\"\".strip()\n",
    "        else:\n",
    "            raise ValueError(f\"Dataset name {dataset_name} not supported\")\n",
    "    \n",
    "        chunk_info += song_info + \"\\n\"\n",
    "\n",
    "    batch_file.append({\n",
    "        \"custom_id\": \"meta-\" + meta[\"id\"],\n",
    "        \"method\": \"POST\",\n",
    "        \"url\": \"/v1/chat/completions\",\n",
    "        \"body\": {\n",
    "            \"model\": \"gpt-4.1\",\n",
    "            \"messages\": [{\"role\": \"system\", \"content\": system_prompt}, {\"role\": \"user\", \"content\": chunk_info}],\n",
    "            \"max_tokens\": 4096\n",
    "        }\n",
    "    })\n",
    "\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_file[0][\"body\"][\"messages\"][1][\"content\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(batch_file)\n",
    "batch_filepath = \"/home/christian/code/christian/metadata/tagging/batchinput_t8.jsonl\"\n",
    "# save the batch file to a jsonl file\n",
    "with open(batch_filepath, \"w\") as f:\n",
    "    for item in batch_file:\n",
    "        f.write(json.dumps(item) + \"\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload the batch file to the client\n",
    "batch_input_file = client.files.create(\n",
    "    file=open(batch_filepath, \"rb\"),\n",
    "    purpose=\"batch\"\n",
    ")\n",
    "\n",
    "print(batch_input_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a batch job\n",
    "batch_input_file_id = batch_input_file.id\n",
    "client.batches.create(\n",
    "    input_file_id=batch_input_file_id,\n",
    "    endpoint=\"/v1/chat/completions\",\n",
    "    completion_window=\"24h\",\n",
    "    metadata={\n",
    "        \"description\": \"batch gpt-4.1 audio tagging genius (10)\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(batch_input_file_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = client.batches.retrieve(\"batch_681175ae91dc819086cd6d074d182d01\")\n",
    "print(batch.request_counts)\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve the batch job\n",
    "file_response = client.files.content(batch.output_file_id)\n",
    "#print(file_response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse the JSONL string into a list of dictionaries\n",
    "responses = [json.loads(line) for line in file_response.text.split('\\n') if line.strip()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for response in tqdm(responses):\n",
    "    item_id = response[\"custom_id\"].replace(\"meta-\", \"\")\n",
    "    content = response[\"response\"][\"body\"][\"choices\"][0][\"message\"][\"content\"]\n",
    "    # parse the content as json\n",
    "    content = json.loads(content)\n",
    "    # get the gpt_tag\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file_response is a string, we need to parse it as jsonl \n",
    "import json\n",
    "from tqdm import tqdm\n",
    "\n",
    "dataset_name = \"imslp\"\n",
    "\n",
    "# Parse the JSONL string into a list of dictionaries\n",
    "responses = [json.loads(line) for line in file_response.text.split('\\n') if line.strip()]\n",
    "\n",
    "results_file = f'/home/christian/code/christian/metadata/tagging/gpt-4_1-audio-tagging_results_{dataset_name}.json'\n",
    "\n",
    "# Initialize results dictionary\n",
    "results = {}\n",
    "\n",
    "# Load existing results if file exists\n",
    "if os.path.exists(results_file):\n",
    "    with open(results_file, 'r') as f:\n",
    "        results = json.load(f)\n",
    "    print(f\"Loaded {len(results)} existing results from {results_file}\")\n",
    "\n",
    "for response in tqdm(responses):\n",
    "    item_id = response[\"custom_id\"].replace(\"meta-\", \"\")\n",
    "    content = response[\"response\"][\"body\"][\"choices\"][0][\"message\"][\"content\"]\n",
    "    \n",
    "    # Add to dictionary with item_id as key\n",
    "    results[item_id] = {\n",
    "        \"gpt_tag\": content\n",
    "    }\n",
    "    \n",
    "    # Also print for verification\n",
    "    #print(f\"ID: {item_id}\")\n",
    "    #print(f\"Content: {content}\")\n",
    "    #print()\n",
    "\n",
    "print(f\"total results: {len(results)}\")\n",
    "\n",
    "# Write the results to a normal JSON file\n",
    "with open(results_file, 'w') as f:\n",
    "    json.dump(results, f, indent=2)\n",
    "\n",
    "print(f\"Results written to {results_file}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Merge with metas\n",
    "\n",
    "This is used to take the gpt tags from json and then merge them into existing training metas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the gpt tags\n",
    "new_tags = {\n",
    "    \"genius\": {},\n",
    "    \"discogs_subset\": {}\n",
    "}\n",
    "for dataset_name in new_tags.keys():\n",
    "    results_file = f'/home/christian/code/christian/metadata/tagging/gpt-4_1-tagging_results_{dataset_name}.json'\n",
    "    gpt_tags = json.load(open(results_file))\n",
    "    print(f\"len(gpt_tags) in new_tags[{dataset_name}]: {len(gpt_tags)}\")\n",
    "    # merge the gpt tags into the metas\n",
    "    new_tags[dataset_name] = gpt_tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get all lond dataset names\n",
    "dataset_names = set([meta[\"dataset\"] for meta in metas])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for subset in [\"tr\"]:\n",
    "    # load the existing metas\n",
    "    base_dir = \"/app/suno/data/chirp_v5_ft/v2/\"\n",
    "    metas = read_jsonl(os.path.join(base_dir, f\"metas_{subset}.jsonl\"))\n",
    "    print(len(metas))\n",
    "\n",
    "    new_metas = []\n",
    "    count = 0\n",
    "    for meta in tqdm(metas):\n",
    "        meta_id = meta[\"id\"]\n",
    "        dataset_name = meta[\"dataset\"]\n",
    "        if \"genius\" in dataset_name:\n",
    "            short_dataset_name = \"genius\"\n",
    "        elif \"discogs_subset\" in dataset_name:\n",
    "            short_dataset_name = \"discogs_subset\"\n",
    "\n",
    "        new_meta = meta.copy()\n",
    "\n",
    "        if short_dataset_name in new_tags:\n",
    "            # check if the meta_id is in the new_tags[dataset_name]\n",
    "            if meta_id in new_tags[short_dataset_name]:\n",
    "                count += 1\n",
    "                # check if we have tags\n",
    "                tags = new_meta.get(\"tags\", [])\n",
    "                gpt_tags = new_tags[short_dataset_name][meta_id][\"gpt_tag\"].split(\";\")\n",
    "                gpt_tags = [tag.strip() for tag in gpt_tags if tag.strip()]\n",
    "                # extend the tags with the gpt tags\n",
    "                tags.extend(gpt_tags)\n",
    "                new_meta[\"tags\"] = tags\n",
    "        #else:\n",
    "        #    meta[\"gpt_tags\"] = \"None\"\n",
    "        new_metas.append(new_meta)\n",
    "    print(f\"count: {count}/{len(metas)}\")\n",
    "    print(len(new_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the new metas to a jsonl file\n",
    "write_jsonl(new_metas, os.path.join(base_dir, f\"metas_{subset}_gpt_tags.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for meta in new_metas:\n",
    "    if \"genius\" in meta[\"dataset\"]:\n",
    "        if meta[\"id\"] in new_tags[\"genius\"]:\n",
    "            print(meta[\"tags\"])\n",
    "            print(len(meta[\"tags\"]))\n",
    "            count += 1\n",
    "\n",
    "    if count > 10:\n",
    "        break\n",
    "print(f\"count: {count}/{len(new_metas)}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read new metas\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "import os\n",
    "subset = \"tr\"\n",
    "base_dir = \"/app/suno/data/chirp_v5_ft/v2/\"\n",
    "\n",
    "\n",
    "old_metas = read_jsonl(os.path.join(base_dir, f\"metas_{subset}.jsonl\"))\n",
    "print(len(old_metas))\n",
    "test_metas = read_jsonl(os.path.join(base_dir, f\"metas_{subset}_gpt_tags.jsonl\"))\n",
    "print(len(test_metas))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_tag_counts = []\n",
    "test_tag_counts = []\n",
    "\n",
    "for old_meta in old_metas:\n",
    "    n_tags = len(old_meta.get(\"tags\", []))\n",
    "    old_tag_counts.append(n_tags)\n",
    "\n",
    "for test_meta in test_metas:\n",
    "    n_tags = len(test_meta.get(\"tags\", []))\n",
    "    test_tag_counts.append(n_tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "total_old_tags_counts = sum(old_tag_counts)\n",
    "total_test_tag_counts = sum(test_tag_counts)\n",
    "print(total_old_tags_counts,  total_test_tag_counts)\n",
    "print(total_test_tag_counts - total_old_tags_counts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Discogs tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we load the original metas to get the artists and song titles\n",
    "METAS_DIR = \"/app/suno/tmp\"\n",
    "raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_discogs_subset_metas.jsonl\"))\n",
    "#raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_genius_metas.jsonl\"))\n",
    "#raw_metas = read_jsonl(os.path.join(METAS_DIR, \"raw_imslp_metas.jsonl\"))\n",
    "print(len(raw_metas))\n",
    "#dataset_name = \"imslp\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_metas = read_jsonl(os.path.join(METAS_DIR, \"clean_discogs_subset_v0_metas.jsonl\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "clean_metas[100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dict mapping from tag to ids and song titles that contain the tag\n",
    "tag_to_ids = {}\n",
    "for meta in raw_metas:\n",
    "    tags = meta.get(\"rym_genres\", [])\n",
    "    for tag in tags:\n",
    "        if tag not in tag_to_ids:\n",
    "            tag_to_ids[tag] = []\n",
    "        tag_to_ids[tag].append(meta)\n",
    "\n",
    "print(f\"len(tag_to_ids): {len(tag_to_ids)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dict mapping from tag to ids and song titles that contain the tag\n",
    "tag_to_ids = {}\n",
    "for meta in clean_metas:\n",
    "    tags = meta.get(\"tags\", [])\n",
    "    for tag in tags:\n",
    "        if tag not in tag_to_ids:\n",
    "            tag_to_ids[tag] = []\n",
    "        tag_to_ids[tag].append(meta)\n",
    "\n",
    "print(f\"len(tag_to_ids): {len(tag_to_ids)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now sort the tag_to_ids by the length of the key (tag)\n",
    "tag_to_ids = sorted(tag_to_ids.items(), key=lambda x: len(x[0]), reverse=False)\n",
    "print(len(tag_to_ids))\n",
    "# get the top 1000 tags\n",
    "for tag, ids in tag_to_ids[:10]:\n",
    "    print(tag, len(ids))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = '/home/christian/code/christian/metadata/discogs_subset_unique_tags.txt'\n",
    "with open(filepath, 'r') as f:\n",
    "    discogs_tags = f.read().splitlines()\n",
    "print(len(discogs_tags))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sort by tag length\n",
    "discogs_tags = sorted(discogs_tags, key=len, reverse=False)\n",
    "for tag in discogs_tags[:50]:\n",
    "    print(tag)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Title cleanning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
