{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"HF_HOME\"] = \"/mnt/localdisk/christian/hf_cache\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4,5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/christian/miniconda3/envs/suno_qwen/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "#from peft import PeftModel\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from accelerate import init_empty_weights, load_checkpoint_and_dispatch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 30/30 [02:57<00:00,  5.91s/it]\n",
      "                                                                                                      \r"
     ]
    }
   ],
   "source": [
    "trained_model_path = \"/home/m4burns/glockenspiel/suno_utils/notebooks/torchtune/llama3_3_70B/lora_lyrics_3/epoch_0\"\n",
    "\n",
    "# Define the model and adapter paths\n",
    "original_model_name = \"/home/christian/llama/Llama-3.3-70B-Instruct\"\n",
    "\n",
    "# Load base model with empty weights first\n",
    "#with init_empty_weights():\n",
    "model = AutoModelForCausalLM.from_pretrained(original_model_name)\n",
    "\n",
    "# Load the model weights across GPUs\n",
    "model = load_checkpoint_and_dispatch(\n",
    "    model,\n",
    "    original_model_name,\n",
    "    device_map=\"auto\",\n",
    "    no_split_module_classes=[\"LlamaDecoderLayer\"],\n",
    "    max_memory={\n",
    "        0: \"70GiB\",\n",
    "        1: \"70GiB\",\n",
    "    },\n",
    "    dtype=\"bfloat16\"\n",
    ")\n",
    "\n",
    "# Load PEFT adapter\n",
    "# NOTE: this will modify the base model in place!!\n",
    "# Calling model.generate after this will use the finetuned model\n",
    "#peft_model = PeftModel.from_pretrained(model, trained_model_path)\n",
    "\n",
    "# Load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(original_model_name)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a music expert. \n",
    "When prompted with an artist or song, \n",
    "you will respond with relevant information releated to the following aspects:\n",
    "\n",
    "{\n",
    "  \"tempo\": \"integer (e.g. 120, 85, 160) - measured in BPM\",\n",
    "  \"key\": \"string (e.g. 'C major', 'F# minor', 'Eb major')\",\n",
    "  \"time_signature\": \"string (e.g. '4/4', '6/8', '3/4', '7/8')\",\n",
    "  \"year\": \"integer (e.g. 1985, 2002, 2024)\",\n",
    "  \"instrumental\": \"boolean (true if no vocals, false if has vocals)\",\n",
    "  \"genre\": [\"array of genres (e.g. ['indie rock', 'shoegaze', 'dream pop'])\", \n",
    "            \"['jazz fusion', 'progressive rock']\",\n",
    "            \"['hip hop', 'trap', 'cloud rap']\"],\n",
    "  \"chord_progression\": \"string (e.g. 'I-IV-V blues progression', 'ii-V-I jazz progression', 'vi-IV-I-V pop progression')\",\n",
    "  \"key_changes\": [\"array of changes (e.g. ['verse in C major', 'chorus modulates to E major'])\",\n",
    "                  \"['starts in D minor', 'bridge shifts to relative major F']\"],\n",
    "  \"song_structure\": \"string (e.g. 'verse-chorus-verse-chorus-bridge-chorus', 'AABA form', 'through-composed')\",\n",
    "  \"instruments\": [\"array (e.g. ['electric guitar', 'bass', 'drums', 'synthesizer'])\",\n",
    "                  \"['piano', 'upright bass', 'trumpet', 'saxophone']\"],\n",
    "  \"specific_gear\": [\"array (e.g. ['Moog Minimoog', 'Roland TR-808', 'Fender Stratocaster'])\",\n",
    "                    \"['Native Instruments Massive', 'SSL mixing console']\"],\n",
    "  \"hook_description\": \"string (e.g. 'descending synth arpeggio', 'four-chord guitar riff', 'repeating bass groove')\",\n",
    "  \"recording_location\": \"string or null (e.g. 'home studio', 'professional studio', 'live venue recording')\",\n",
    "  \"production_style\": \"string (e.g. 'lo-fi', 'polished studio production', 'raw garage recording', 'minimalist')\",\n",
    "  \"audio_effects\": [\"array (e.g. ['reverb', 'delay', 'distortion'])\",\n",
    "                    \"['auto-tune', 'compression', 'chorus']\"],\n",
    "  \"vocal_style\": \"string or null (e.g. 'whispered', 'belting', 'falsetto', 'spoken word', 'harmonized')\",\n",
    "  \"vocalist_gender\": \"string or null (e.g. 'male', 'female', 'multiple', null for instrumental)\",\n",
    "  \"dynamics\": \"string (e.g. 'quiet verses with loud chorus', 'consistently loud', 'gradually building')\",\n",
    "  \"melody_characteristics\": \"string (e.g. 'pentatonic scale runs', 'chromatic passages', 'simple repetitive motif')\",\n",
    "  \"lyrical_themes\": [\"array (e.g. ['love', 'loss', 'social justice'])\",\n",
    "                    \"['nature', 'spirituality', 'personal growth']\"],\n",
    "  \"typical_usage\": [\"array (e.g. ['workout', 'meditation', 'party'])\",\n",
    "                    \"['film soundtrack', 'commercial background', 'dance']\"],\n",
    "  \"chart_performance\": \"string or null (e.g. 'peaked at #5', 'charted in multiple countries', null if never charted)\",\n",
    "  \"cultural_impact\": \"string (e.g. 'influenced future genre development', 'widely sampled', 'cult following')\",\n",
    "  \"critical_reception\": \"string (e.g. 'universally acclaimed', 'mixed reviews', 'underground favorite')\",\n",
    "  \"album_artwork\": \"string or null (e.g. 'minimalist black and white photograph', 'abstract digital art', 'hand-drawn illustration')\",\n",
    "  \"music_video\": \"string or null (e.g. 'performance-based', 'narrative storyline', 'abstract visuals')\"\n",
    "}\n",
    "\n",
    "Be as concise, but provide details where noteable or differenting. \n",
    "Ensure to not reference any artists, albums, or names of individuals. \n",
    "Your response should be formatted in valid JSON. \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "chat = [\n",
    "  {\"role\": \"system\", \"content\": system_prompt},\n",
    "  {\"role\": \"user\", \"content\": \"Need 2 by Pinegrove\"},\n",
    "]\n",
    "\n",
    "toks = tokenizer.apply_chat_template(chat, tokenize=True, add_bos=True, return_tensors=\"pt\").to(\"cuda\")\n",
    "outputs = model.generate(toks, max_new_tokens=2048)\n",
    "output_text = tokenizer.decode(outputs[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract json from output_json\n",
    "import json\n",
    "\n",
    "try:\n",
    "    json_str = output_text.split(\"```\")[1].strip()\n",
    "    data_dict = json.loads(json_str)\n",
    "except:\n",
    "    try:\n",
    "        json_str = output_text.split(\"```json\")[1].strip(\"```<|eot_id|>\")\n",
    "        data_dict = json.loads(json_str)\n",
    "    except:\n",
    "        print(\"error\")\n",
    "        #print(output_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tempo: 96\n",
      "duration: 3:18\n",
      "key: G major\n",
      "time_signature: 4/4\n",
      "year: 2016\n",
      "instrumental: False\n",
      "genre: ['indie rock', 'emo', 'folk punk']\n",
      "chord_progression: I-V-vi-IV progression with variations\n",
      "key_changes: ['verse and chorus in G major', 'bridge modulates to Em']\n",
      "song_structure: verse-chorus-verse-chorus-bridge-chorus\n",
      "instruments: ['acoustic guitar', 'electric guitar', 'bass', 'drums']\n",
      "specific_gear: ['Fender Telecaster', ' Gibson Les Paul']\n",
      "hook_description: repeating guitar riff and vocal melody\n",
      "recording_location: home studio\n",
      "production_style: lo-fi with polished elements\n",
      "audio_effects: ['reverb', 'delay']\n",
      "vocal_style: heartfelt, emotive delivery\n",
      "vocalist_gender: male\n",
      "dynamics: quiet verses with louder choruses\n",
      "melody_characteristics: simple, memorable motifs\n",
      "lyrical_themes: ['self-discovery', ' relationships', 'growth']\n",
      "typical_usage: ['study', 'relaxation', 'focus']\n",
      "chart_performance: None\n",
      "cultural_impact: influenced the emo revival movement\n",
      "critical_reception: widely acclaimed for its emotional intensity\n",
      "album_artwork: watercolor painting of a natural landscape\n",
      "music_video: lyric-based animation with natural imagery\n"
     ]
    }
   ],
   "source": [
    "for key, val in data_dict.items():\n",
    "    print(f\"{key}: {val}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_qwen",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}