{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2e2253c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dabe64d3-8bc5-444c-8f94-b54167ec015f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-22T03:21:36.087854Z",
     "iopub.status.busy": "2024-10-22T03:21:36.087174Z",
     "iopub.status.idle": "2024-10-22T03:21:36.677530Z",
     "shell.execute_reply": "2024-10-22T03:21:36.677128Z",
     "shell.execute_reply.started": "2024-10-22T03:21:36.087823Z"
    }
   },
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "client = OpenAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ac46fe09",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_artists = \"\"\"\n",
    "Aitana\n",
    "Aitch\n",
    "Ajay Chandisar\n",
    "Ajay Gogavale\n",
    "Ajay Mali\n",
    "Ajay Thakor\n",
    "Ajay-Atul\n",
    "Ajda Pekkan\n",
    "Ajeet\n",
    "Ajeet Anand\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "dfb11d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_response(input_artists):\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": f\"\"\"\n",
    "          You are a music expert.\n",
    "          Given a list of input artists, please annotate each artist's music style.\n",
    "          Try to be as accurate as possible, but also comprehensive as possible.\n",
    "          Only annotate the artists that are in the input list.\n",
    "          Keep the answer to be max 7 different styles per artist.\n",
    "          Cross reference with your knowledge of music to make sure the answer is correct.\n",
    "          Remove all the non-confident answers please.\n",
    "        \n",
    "          Format the output as a JSON object where the keys are the artist names and the values are the styles.\n",
    "          The values should be a comma-separated list of styles.\n",
    "\n",
    "          \n",
    "          Here is the list of artists (separated by new lines):\n",
    "          \"{input_artists}\"\n",
    "          \"\"\",\n",
    "                    }\n",
    "                ],\n",
    "            },\n",
    "        ],\n",
    "        temperature=1,\n",
    "        max_tokens=4096,\n",
    "        top_p=1,\n",
    "        frequency_penalty=0,\n",
    "        presence_penalty=0,\n",
    "        response_format={\"type\": \"text\"},\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "193b48ec-1ec0-454c-ad1f-a1746a949eff",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.chat.completions.create(\n",
    "    model=\"gpt-4o\",\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": f\"\"\"\n",
    "          You are a music expert.\n",
    "          Given a list of input artists, please annotate each artist's music style.\n",
    "          Try to be as accurate as possible, but also comprehensive as possible.\n",
    "          Only annotate the artists that are in the input list.\n",
    "          Keep the answer to be max 7 different styles per artist.\n",
    "          Cross reference with your knowledge of music to make sure the answer is correct.\n",
    "          Remove all the non-confident answers please.\n",
    "        \n",
    "          Format the output as a JSON object where the keys are the artist names and the values are the styles.\n",
    "          The values should be a comma-separated list of styles.\n",
    "\n",
    "          \n",
    "          Here is the list of artists (separated by new lines):\n",
    "          \"{input_artists}\"\n",
    "          \"\"\",\n",
    "                }\n",
    "            ],\n",
    "        },\n",
    "    ],\n",
    "    temperature=1,\n",
    "    max_tokens=4096,\n",
    "    top_p=1,\n",
    "    frequency_penalty=0,\n",
    "    presence_penalty=0,\n",
    "    response_format={\"type\": \"text\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "4bc5e0e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21238\n"
     ]
    }
   ],
   "source": [
    "with open(\n",
    "    \"/Users/tonytong/Work/Git/glockenspiel/suno_utils/suno_utils/worker/assets/top_artists.txt.py\"\n",
    ") as f:\n",
    "    TOP_ARTISTS = [line.strip() for line in f.readlines()]\n",
    "\n",
    "print(len(TOP_ARTISTS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "52ae54bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from collections import defaultdict\n",
    "\n",
    "total_ans = defaultdict(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "8e58c328",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "\n",
    "# # Remove 'json\\n' if present, and any leading/trailing backticks and newlines\n",
    "# cleaned_ans = response.choices[0].message.content.replace(\"json\\n\", \"\").strip(\"`\\n\")\n",
    "\n",
    "# # Parse the cleaned string into a dictionary\n",
    "# curr_ans = json.loads(cleaned_ans)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4d31ef",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "c786577a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "b2e67da6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21238/21238 [00:00<00:00, 3214094.47it/s]\n"
     ]
    }
   ],
   "source": [
    "total_ans = defaultdict(str)\n",
    "for artist in tqdm.tqdm(TOP_ARTISTS):\n",
    "    total_ans[artist] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "0a13dc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in tqdm.tqdm(range(0, len(TOP_ARTISTS), 20)):\n",
    "for i in [2160, 3280, 3900, 4080, 5280, 8980, 10920, 15460, 18280, 19260]:\n",
    "    artist_chunk = \"\\n\".join(TOP_ARTISTS[i : i + 20])\n",
    "    response = get_response(artist_chunk)\n",
    "    try:\n",
    "        curr_ans = (\n",
    "            response.choices[0].message.content.replace(\"json\\n\", \"\").strip(\"`\\n\")\n",
    "        )\n",
    "        curr_ans = json.loads(curr_ans)\n",
    "        if isinstance(curr_ans, dict):\n",
    "            for artist, generes in curr_ans.items():\n",
    "                total_ans[artist] = generes\n",
    "    except Exception as e:\n",
    "        print(f\"Error on chunk {i}, {e}\")\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "e147bc77",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "44 ['Adana Zapata', 'Agata Christie', 'Arsın Uzunov', 'Cheikh Niang', 'Duzu Aharon', 'Eizm', 'Empire of the Sun', 'Hamza El Shaeri', 'Jordyn Jones', 'Junior Kelly Marchena', 'K/I.Z', 'Kavva', 'Mary G J. Blige', 'MusikalBasics', 'Nababález', 'Niloofur Usmonova', 'Pablo Milanés', 'Rinko Gaetano', 'Sangeeta', 'Sarikodie', 'Shinsei Kamattechan', 'Timothai', 'Trick', 'Yuzu', 'Yandeh', 'ZASTEROKO La Resistencia Salsera del Callao', 'Șaban Gürsoy', 'Șaníșer', 'Șebnem Ferah', 'Șebnem Tovuzlu', 'Șehinșah', 'Șevval Sam', 'Șivan Perwer', 'Șiyar Berwari', 'Șiyar û Dijwar', 'Șöhret Memmedov', 'Șaban Șaulić', 'Șako Polumenta', 'Șerif Konjević', 'Șkwor', 'Ștístko a Poupěnka', 'bbno$', 'iLLEOo', 'rusina']\n"
     ]
    }
   ],
   "source": [
    "mistaken_artists = []\n",
    "unknown_artists = set(TOP_ARTISTS)\n",
    "for k in total_ans:\n",
    "    if k not in unknown_artists:\n",
    "        mistaken_artists.append(k)\n",
    "print(len(mistaken_artists), mistaken_artists)\n",
    "for k in mistaken_artists:\n",
    "    total_ans.pop(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "685987bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_ans = {}\n",
    "for artist in TOP_ARTISTS:\n",
    "    if artist in total_ans:\n",
    "        final_ans[artist] = total_ans[artist]\n",
    "    else:\n",
    "        final_ans[artist] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "b30cec7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"final_ans.json\", \"w\") as f:\n",
    "    json.dump(final_ans, f, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "7f694ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_response_second_pass(input_artists):\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": f\"\"\"\n",
    "          You are a music expert with strong domain knowledge.\n",
    "          Given a list of input artists and their guessed music styles, \n",
    "          please annotate or translate each artist's music style into only English words.\n",
    "          Try to be as accurate as possible, correct the wrong styles, but add the missing styles.\n",
    "          Only annotate the artists that are in the input list.\n",
    "          Keep the answer to be max 7 different styles per artist.\n",
    "          Cross reference with your knowledge of music to make sure the answer is correct.\n",
    "          Remove all the non-confident styles please.\n",
    "        \n",
    "          Format the output as a JSON object where the keys are the artist names and the values are the styles.\n",
    "          The values should be a comma-separated list of styles.\n",
    "          \n",
    "          Artist name is followed by \":\", then the guessed styles.\n",
    "          Like \"ArtistA: style1, style2, style3\"\n",
    "          Each artist is separated by a new line.\n",
    "          Here is the list of artists and their styles.\n",
    "          \"{input_artists}\"\n",
    "          \"\"\",\n",
    "                    }\n",
    "                ],\n",
    "            },\n",
    "        ],\n",
    "        temperature=1,\n",
    "        max_tokens=4096,\n",
    "        top_p=1,\n",
    "        frequency_penalty=0,\n",
    "        presence_penalty=0,\n",
    "        response_format={\"type\": \"text\"},\n",
    "    )\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "efd4ca46",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_ans = defaultdict(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "5a117f0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1062/1062 [1:02:49<00:00,  3.55s/it]\n"
     ]
    }
   ],
   "source": [
    "missed_indices = []\n",
    "for i in tqdm.tqdm(range(0, len(TOP_ARTISTS), 20)):\n",
    "    artist_chunk = \"\"\n",
    "    for artist in TOP_ARTISTS[i : i + 20]:\n",
    "        artist_chunk += f\"{artist}: {final_ans[artist]}\\n\"\n",
    "    response = get_response_second_pass(artist_chunk)\n",
    "    try:\n",
    "        curr_ans = (\n",
    "            response.choices[0].message.content.replace(\"json\\n\", \"\").strip(\"`\\n\")\n",
    "        )\n",
    "        curr_ans = json.loads(curr_ans)\n",
    "        if isinstance(curr_ans, dict):\n",
    "            for artist, generes in curr_ans.items():\n",
    "                total_ans[artist] = generes\n",
    "    except Exception as e:\n",
    "        print(f\"Error on chunk {i}, {e}\")\n",
    "        missed_indices.append(i)\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "9f7f4e2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_ans = {}\n",
    "for artist in TOP_ARTISTS:\n",
    "    if artist in total_ans:\n",
    "        final_ans[artist] = total_ans[artist]\n",
    "    else:\n",
    "        final_ans[artist] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "fca7935a",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"final_ans_second_pass.json\", \"w\") as f:\n",
    "    json.dump(final_ans, f, indent=4, ensure_ascii=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "1813ec2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21238\n"
     ]
    }
   ],
   "source": [
    "with open(\n",
    "    \"/Users/tonytong/Work/Git/glockenspiel/suno_utils/suno_utils/worker/assets/top_artists.txt.py\"\n",
    ") as f:\n",
    "    TOP_ARTISTS = [line.strip() for line in f.readlines()]\n",
    "\n",
    "print(len(TOP_ARTISTS))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "013c6edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for top_artist in TOP_ARTISTS:\n",
    "    if top_artist not in final_ans:\n",
    "        print(top_artist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d72632e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
