{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "\n",
    "import json\n",
    "import funcy\n",
    "import torch\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from suno_utils.utils.text import read_jsonl\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl\", progress=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "rand_meta = metas[rand_idx]\n",
    "#for key in rand_meta:\n",
    "youtube_title = rand_meta[\"youtube_title\"]\n",
    "tags = rand_meta[\"tags_text\"]\n",
    "print(youtube_title, tags)\n",
    "\n",
    "# download from s3\n",
    "s3_filepath = rand_meta[\"audio_filepath\"]\n",
    "# download from s3 \n",
    "filename = os.path.basename(s3_filepath)\n",
    "print(filename)\n",
    "os.system(f\"\"\"aws s3 cp {s3_filepath} /home/christian/code/christian/notebooks/outputs/s3\"\"\")\n",
    "local_path = f\"/home/christian/code/christian/notebooks/outputs/s3/{filename}\"\n",
    "print(local_path)\n",
    "audio, sr = torchaudio.load(local_path)\n",
    "if audio.abs().max() > 1.0:\n",
    "    audio = audio / audio.abs().max()\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(data=audio, rate=sr, normalize=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# count the number of tags in each song\n",
    "tag_counts_per_meta = []\n",
    "for meta in metas:\n",
    "    tags = meta[\"tags_text\"]\n",
    "    tag_counts_per_meta.append(len(tags))\n",
    "\n",
    "# make a histogram of the tag counts\n",
    "tag_counts_per_meta = np.array(tag_counts_per_meta)\n",
    "import matplotlib.pyplot as plt\n",
    "plt.hist(tag_counts_per_meta, bins=30)\n",
    "print(\"min\", np.min(tag_counts_per_meta), \"max\", np.max(tag_counts_per_meta), \"mean\", np.mean(tag_counts_per_meta), \"median\", np.median(tag_counts_per_meta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "When provided with the title of a music video you respond with a comma separated single list of descriptors. \n",
    "Make sure to include detailed genre tags, mood, style, time period, tempo, and vocal type (gender, delivery, and processing). \n",
    "If notable, also include details of production, audio gear (e.g. mic type) and instruments (e.g. guitar amp type) used.\n",
    "If the song has lyrics include the language. Do not include reference to any musicians or musical artists. \n",
    "Include as many descriptors as possible.\n",
    "\"\"\"\n",
    "\n",
    "# 1) genre, 2) mood, 3) production techniques, 4) audio equipment used, 5) instruments used, 6) time period (year),  7) tempo, 8) male or female vocal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "\n",
    "#client = OpenAI(api_key=\"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\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "rand_meta = metas[rand_idx]\n",
    "#for key in rand_meta:\n",
    "youtube_title = rand_meta[\"youtube_title\"]\n",
    "print(youtube_title)\n",
    "\n",
    "completion = client.chat.completions.create(\n",
    "  model=\"gpt-4o\",\n",
    "  messages=[\n",
    "    {\"role\": \"system\", \"content\": system_prompt},\n",
    "    {\"role\": \"user\", \"content\": f\"{youtube_title}\"}\n",
    "  ]\n",
    ")\n",
    "\n",
    "string = completion.choices[0].message.content\n",
    "\n",
    "print(string)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "string = completion.choices[0].message.content\n",
    "tags = string.split(\",\")\n",
    "tags = [tag.strip() for tag in tags]\n",
    "tags = [tag.strip(\".\") for tag in tags]\n",
    "for tag in tags:\n",
    "    print(tag.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# here we loop through all the metas and generate tags for each song\n",
    "output_metas_map_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas_filtered_enhanced.json\"\n",
    "# use dict by dataset item id to store the tags\n",
    "\n",
    "\n",
    "if os.path.exists(output_metas_map_filepath):\n",
    "    with open(output_metas_map_filepath, \"r\") as f:\n",
    "        output_metas_map = json.load(f)\n",
    "    print(\"loaded existing tags: \", len(output_metas_map))\n",
    "else:\n",
    "    output_metas_map = {}\n",
    "\n",
    "max_examples = 60\n",
    "#max_examples = len(metas)\n",
    "\n",
    "for n in tqdm(range(max_examples)):\n",
    "    #n = np.random.randint(0, len(metas))\n",
    "    meta = metas[n]\n",
    "    meta_id = meta[\"id\"]\n",
    "    if meta_id in output_metas_map:\n",
    "        continue\n",
    "\n",
    "    youtube_title = meta[\"youtube_title\"]\n",
    "    completion = client.chat.completions.create(\n",
    "      model=\"gpt-4o\",\n",
    "      messages=[\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": f\"{youtube_title}\"}\n",
    "      ]\n",
    "    )\n",
    "    string = completion.choices[0].message.content\n",
    "    tags = string.split(\",\")\n",
    "    tags = [tag.strip() for tag in tags]\n",
    "    tags = [tag.strip(\".\") for tag in tags]\n",
    "    output_metas_map[meta_id] = {\"tags\" : tags, \"youtube_title\": youtube_title}\n",
    "\n",
    "    # save to disk\n",
    "    with open(output_metas_map_filepath, \"w\") as f:\n",
    "        json.dump(output_metas_map, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(output_metas_map_filepath, \"r\") as f:\n",
    "    output_metas_map = json.load(f)\n",
    "\n",
    "print(\"loaded existing tags: \", len(output_metas_map))\n",
    "\n",
    "for meta_id, val in output_metas_map.items():\n",
    "    tags = val[\"tags\"]\n",
    "    youtube_title = val[\"youtube_title\"]\n",
    "    print(youtube_title)\n",
    "    print(tags)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mistral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "device = \"cuda\" # the device to load the model onto\n",
    "access_token = \"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\"\n",
    "model = AutoModelForCausalLM.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\", token=access_token)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"mistralai/Mistral-7B-Instruct-v0.1\", token=access_token)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import InferenceClient\n",
    "\n",
    "client = InferenceClient(\n",
    "    \"mistralai/Mistral-7B-Instruct-v0.1\",\n",
    "    token=access_token,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mistral_system_prompt = \"\"\"\n",
    "Describe the song with a comma separated list of words. These words should include detailed genre tags, mood, style, time period, tempo, and vocal type. \n",
    "If notable, also include details of production, audio gear (e.g. mic type) and instruments (e.g. guitar amp type) used.\n",
    "If the song has lyrics include the language. Do not include reference to any musicians or musical artists. \n",
    "Do not respond with multiple sentences. Do not repeat the title of the song.\n",
    "Include as much detail as possible but make sure your entire response is a single list of words. \n",
    "The title of the video is: \n",
    "\"\"\"\n",
    "\n",
    "mistral_example_title_1 = \"Party in LA - Miley Cyrus\"\n",
    "mistral_example_response_1 = \"Pop, 2010s, upbeat, carefree, modern, mid-tempo, English lyrics, catchy hooks, polished production, digital recording, bright synths, groovy bassline, clean electric guitar, danceable rhythm, youthful energy, fun mood, vibrant, radio-friendly, layered vocals, punchy drums, contemporary pop, light reverb, compressed dynamics, summer vibe, party atmosphere.\"\n",
    "mistral_example_title_2 = \"Dreams - Fleetwood Mac\"\n",
    "mistral_example_response_2 = \"Soft rock, 1970s, mellow, laid-back, melancholic, moderate tempo, English lyrics, smooth vocals, warm production, analog recording, Fender Rhodes, clean electric guitar, light percussion, natural reverb, vintage vibe, introspective mood, subtle harmonies, analog tape warmth, dynamic range, emotional tone, intimate feel, soft sustain, understated bass, acoustic drum kit.\"\n",
    "mistral_example_title_3 = \"Enter Sandman - Metallica\"\n",
    "mistral_example_response_3 = \"Heavy metal, 1990s, dark, aggressive, high energy, fast tempo, English lyrics, distorted guitars, powerful vocals, driving bassline, thunderous drums, analog recording, tight production, deep reverb, punchy dynamics, ominous mood, crunchy guitar riffs, solid-state amps, gritty tone, headbanging rhythm, intense atmosphere, raw power, layered guitars, stadium rock, heavy distortion, suspenseful intro, commanding presence.\"\n",
    "\n",
    "\n",
    "messages = [\n",
    "\t{\"role\" : \"user\", \"content\" : f\"{mistral_system_prompt} {mistral_example_title_1}\"},\n",
    "\t{\"role\" : \"assistant\", \"content\" : mistral_example_response_1},\n",
    "    {\"role\" : \"user\", \"content\" : f\"{mistral_example_title_2}\"},\n",
    "    {\"role\" : \"assistant\", \"content\" : mistral_example_response_2},\n",
    "    {\"role\" : \"user\", \"content\" : f\"{mistral_example_title_3}\"},\n",
    "    {\"role\" : \"assistant\", \"content\" : mistral_example_response_3},\n",
    "\t{\"role\" : \"user\", \"content\" : f\"{youtube_title}\"},\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "Given the youtube title of a song provide a list of the following descriptors in a single JSON object format. \n",
    "The descriptors should include the style, rhythm, key, tempo, singer type (male/female), \n",
    "language of vocals, instruments, audio effects, mood, mixing/mastering terms, usage context, chart performance, \n",
    "and time period. If you are unfamiliar with the song return \"None\" only. Do not return the title or artist. \n",
    "Be as descriptive as possible\"\"\"\n",
    "prompt_song = \"6 Dogs - The Dash (Mindframes) [Official Lyric Video] 6-dogs-the-dash-mindframes-lyrics\"\n",
    "example_response = \"\"\"\n",
    "{\n",
    "  \"style\": \"Cloud rap, trap\",\n",
    "  \"rhythm\": \"Syncopated trap beat with a laid-back groove\",\n",
    "  \"key\": \"Minor key, likely D minor or a similar moody scale\",\n",
    "  \"tempo\": \"Moderate tempo, around 120-140 BPM\",\n",
    "  \"singer_type\": \"Male\",\n",
    "  \"language_of_vocals\": \"English\",\n",
    "  \"instruments\": \"Synth pads, 808 bass, hi-hats, snare rolls, digital kick drum\",\n",
    "  \"audio_effects\": \"Autotune on vocals, reverb, delay, and some stereo imaging effects\",\n",
    "  \"mood\": \"Reflective, melancholic, dreamy\",\n",
    "  \"mixing/mastering_terms\": \"Cleanly mixed vocals, heavy low-end emphasis, crisp high-end on percussions\",\n",
    "  \"usage_context\": \"Casual listening, introspection, chill-out sessions\",\n",
    "  \"chart_performance\": \"None\",\n",
    "  \"time_period\": \"2020s\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "print(f\"{prompt}: {prompt_song}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "#rand_idx = 10001\n",
    "rand_meta = metas[rand_idx]\n",
    "#for key in rand_meta:\n",
    "print(rand_meta)\n",
    "youtube_title = rand_meta[\"youtube_title\"]\n",
    "genius_slug = rand_meta[\"genius_slug\"]\n",
    "tags = \", \".join(rand_meta[\"tags_text\"])\n",
    "print(youtube_title, genius_slug, tags)\n",
    "\n",
    "messages = [\n",
    "\t{\"role\" : \"user\", \"content\" : f\"{prompt}: {prompt_song}\"},\n",
    "\t{\"role\" : \"assistant\", \"content\" : example_response},\n",
    "\t{\"role\" : \"user\", \"content\" : f\"{prompt}: {youtube_title} {genius_slug}\"},\n",
    "]\n",
    "\n",
    "with torch.no_grad():\n",
    "\tencodeds = tokenizer.apply_chat_template(messages, return_tensors=\"pt\")\n",
    "\tmodel_inputs = encodeds.to(device)\n",
    "\tmodel.to(device)\n",
    "\tgenerated_ids = model.generate(model_inputs, max_new_tokens=1024, do_sample=True)\n",
    "\tdecoded = tokenizer.batch_decode(generated_ids)\n",
    "\tresponse = decoded[0].split(\"[/INST]\")[-1]\n",
    "\tresponse = response.strip(\"</s>\")\n",
    "\tresponse = response.strip(\"\\\"\")\n",
    "\tresponse = response.strip(\"[\")\n",
    "\tresponse = response.strip(\"]\")\n",
    "\n",
    "\tprint(response)\n",
    "\tresponse = response.replace(\"\\\\\", \"\")\n",
    "\t# strip anything outside of {}\n",
    "\tresponse = response.split(\"{\")[1].split(\"}\")[0]\n",
    "\t# add {} back\n",
    "\tresponse = \"{\" + response + \"}\"\n",
    "\t# convert to json\n",
    "\tresponse = json.loads(response)\n",
    "\t#print(response)\n",
    "\tfor key, val in response.items():\n",
    "\t\tprint(key, val)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# actual inference\n",
    "# here we loop through all the metas and generate tags for each song\n",
    "output_metas_map_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas_filtered_enhanced_mistral.json\"\n",
    "# use dict by dataset item id to store the tags\n",
    "\n",
    "if os.path.exists(output_metas_map_filepath):\n",
    "    with open(output_metas_map_filepath, \"r\") as f:\n",
    "        output_metas_map = json.load(f)\n",
    "    print(\"loaded existing tags: \", len(output_metas_map))\n",
    "else:\n",
    "    output_metas_map = {}\n",
    "\n",
    "max_examples = 1000\n",
    "#max_examples = len(metas)\n",
    "\n",
    "for n in tqdm(range(max_examples)):\n",
    "    #n = np.random.randint(0, len(metas))\n",
    "    meta = metas[n]\n",
    "    meta_id = meta[\"id\"]\n",
    "    if meta_id in output_metas_map:\n",
    "        continue\n",
    "\n",
    "    youtube_title = meta[\"youtube_title\"]\n",
    "\n",
    "    # run the model\n",
    "    messages = [\n",
    "        {\"role\" : \"user\", \"content\" : f\"{mistral_system_prompt} {mistral_example_title}\"},\n",
    "        {\"role\" : \"assistant\", \"content\" : mistral_example_response},\n",
    "        {\"role\" : \"user\", \"content\" : f\"{mistral_system_prompt} {youtube_title}\"},\n",
    "    ]\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        encodeds = tokenizer.apply_chat_template(messages, return_tensors=\"pt\")\n",
    "        model_inputs = encodeds.to(device)\n",
    "        model.to(device)\n",
    "        generated_ids = model.generate(model_inputs, max_new_tokens=256, do_sample=True)\n",
    "        decoded = tokenizer.batch_decode(generated_ids)\n",
    "        response = decoded[0].split(\"[/INST]\")[-1]\n",
    "        response = response.strip(\"</s>\")\n",
    "        response = response.strip(\"\\\"\")\n",
    "        response = response.strip(\"[\")\n",
    "        response = response.strip(\"]\")\n",
    "        tags = response.split(\",\")\n",
    "        tags = [tag.strip() for tag in tags]\n",
    "        tags = [tag.strip(\"\\\"\") for tag in tags]\n",
    "        tags = [tag.strip(\".\") for tag in tags]\n",
    "\n",
    "    new_tags = []\n",
    "    for tag in tags:\n",
    "        # check for newlines\n",
    "        if \"\\n\" in tag:\n",
    "            sub_tags = tag.split(\"\\n\")\n",
    "            for sub_tag in sub_tags:\n",
    "                new_tags.append(sub_tag)\n",
    "        else:\t\n",
    "            new_tags.append(tag)\n",
    "\n",
    "    output_metas_map[meta_id] = {\"tags\" : new_tags, \"youtube_title\": youtube_title}\n",
    "\n",
    "    # save to disk every 100 iterations\n",
    "    if n % 100 == 0:\n",
    "        with open(output_metas_map_filepath, \"w\") as f:\n",
    "            json.dump(output_metas_map, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(output_metas_map_filepath, \"r\") as f:\n",
    "    output_metas_map = json.load(f)\n",
    "\n",
    "print(\"loaded existing tags: \", len(output_metas_map))\n",
    "\n",
    "for meta_id, val in output_metas_map.items():\n",
    "    tags = val[\"tags\"]\n",
    "    youtube_title = val[\"youtube_title\"]\n",
    "    print(youtube_title)\n",
    "    print(tags)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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
}
