{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import labelbox as lb\n",
    "from datetime import datetime\n",
    "\n",
    "from langdetect import detect\n",
    "from bs4 import BeautifulSoup\n",
    "from suno_utils.utils.s3 import check_s3_file_exists\n",
    "from tqdm import tqdm\n",
    "\n",
    "import tempfile\n",
    "from suno_utils.utils.s3 import upload_s3_files, apply_on_s3, download_s3_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "####### UPDATE THESE ##########\n",
    "models = [\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v4-h-s-32_2025_04_01-14_40_54.json\",\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v5-sem-6b_2025_04_01-15_24_24.json\",\n",
    "    \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v5-sem-6b_repeat_and_neg_2025_04_01-21_25_51.json\",\n",
    "]\n",
    "\n",
    "model_names = [\"v4\", \"auk\", \"auk_aug\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_gen_json(filepath, model_name):\n",
    "    with open(filepath, \"r\", encoding=\"utf-8\") as file:\n",
    "        data = json.load(file)\n",
    "    by_item = []\n",
    "    for genre, outputs in data.items():\n",
    "        for val in outputs:\n",
    "            val[\"genre\"] = genre\n",
    "            val[\"model\"] = model_name\n",
    "            by_item.append(val)\n",
    "            val[\"instrumental\"] = len(val[\"lyrics\"]) < 15\n",
    "            if val[\"instrumental\"]:\n",
    "                val[\"lyrics\"] = [\"Instrumental\"]\n",
    "    df = pd.DataFrame(by_item)\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = [load_gen_json(models[i], model_names[i]) for i in range(0, len(models))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs[2].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_english(df):\n",
    "    filtered = df[df[\"instrumental\"] == False]\n",
    "    for text in filtered.lyrics.tolist():\n",
    "        assert detect(text) == \"en\"\n",
    "\n",
    "\n",
    "for df in dfs:\n",
    "    is_english(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dataset\n",
    "client = lb.Client(\n",
    "    api_key=\"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOiJjbHppbmJnY2wwMDYyMDd5bWg2enhiMTd6Iiwib3JnYW5pemF0aW9uSWQiOiJjbHppbmJnY2QwMDYxMDd5bWM4cjM5cHdzIiwiYXBpS2V5SWQiOiJjbHp3cmU3ZnQwYjFzMDd6aWRrNTFnb21mIiwic2VjcmV0IjoiMGU5M2MwNmU4ZWI1Y2Y3NTlmNTk5YTk5MzIwOTU5MzQiLCJpYXQiOjE3MjM4MTU3NTcsImV4cCI6MjM1NDk2Nzc1N30.Z6gZwlzQ85KrqGOydqCdo1RVmUhOEntpev2HhNso4PU\"\n",
    ")\n",
    "dataset = client.create_dataset(name=\"gender-prompt-adherence-test3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timestamp = datetime.now().strftime(\"%Y%m%d\")\n",
    "test_label = f\"gender-adherence-{timestamp}\"\n",
    "\n",
    "s3_bucket = f\"s3://suno-annotation-public/preference-{test_label}/\"  # this must have trailing slash\n",
    "s3_bucket_url = f\"https://suno-annotation-public.s3.amazonaws.com/preference-{test_label}\"\n",
    "output_html_dir = os.path.abspath(f\"./outputs/html-{test_label}\")\n",
    "os.makedirs(output_html_dir, exist_ok=True)\n",
    "base_html_filepath = \"./templates/singer_gender.html\"\n",
    "global_keys = []\n",
    "\n",
    "# open the html file\n",
    "# Read the HTML file\n",
    "with open(base_html_filepath, \"r\", encoding=\"utf-8\") as file:\n",
    "    soup = BeautifulSoup(file, \"html.parser\")\n",
    "\n",
    "metadata = {}\n",
    "\n",
    "assets = []\n",
    "html_local_paths = []\n",
    "s3_paths = []\n",
    "s3_upload_paths = []\n",
    "html_upload_paths = []\n",
    "for df in dfs:\n",
    "    for i in tqdm(range(0, len(df))):\n",
    "        clip = df.iloc[i]\n",
    "\n",
    "        # use request id as the unique identifier\n",
    "        request_id = f\"{clip.s3_id}_{test_label}\"\n",
    "\n",
    "        orig_pref = clip.vocalist  # by default the second clip is the preferred one\n",
    "\n",
    "        clip_s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{clip.s3_id}.mp3\"\n",
    "        assert check_s3_file_exists(clip_s3_filepath)\n",
    "\n",
    "        # edit the html file to add audios and prompt information\n",
    "        soup.find(id=\"audio1\")[\"src\"] = f\"{s3_bucket_url}/{clip.s3_id}.mp3\"\n",
    "\n",
    "        # Find the elements by ID\n",
    "        lyrics_div = soup.find(id=\"lyrics\")\n",
    "        tags_div = soup.find(id=\"tags\")\n",
    "\n",
    "        # Clear the existing content\n",
    "        lyrics_div.clear()\n",
    "        tags_div.clear()\n",
    "\n",
    "        # Insert the content directly\n",
    "        tags_div.string = f\"Tags: {clip.tags_no_gender}\"\n",
    "        lyrics_div.string = f\"Lyrics: {clip.lyrics}\"\n",
    "\n",
    "        # Save the modified HTML back to disk\n",
    "        output_filepath = os.path.join(output_html_dir, f\"{request_id}.html\")\n",
    "        with open(output_filepath, \"w\", encoding=\"utf-8\") as fp:\n",
    "            fp.write(str(soup))\n",
    "\n",
    "        # push html file to s3 public bucket\n",
    "        s3_mp3_path = f\"{s3_bucket}{clip.s3_id}.mp3\"\n",
    "        s3_html_filepath = f\"{s3_bucket}{request_id}.html\"\n",
    "        s3_html_url = f\"{s3_bucket_url}/{request_id}.html\"\n",
    "        s3_paths.append(clip_s3_filepath)\n",
    "        s3_upload_paths.append(s3_mp3_path)\n",
    "        html_local_paths.append(output_filepath)\n",
    "        html_upload_paths.append(s3_html_filepath)\n",
    "        # os.system(f\"aws s3 cp {output_filepath} {s3_html_filepath}\")\n",
    "\n",
    "        # push mp3 files to s3\n",
    "        # s3_mp3_path = f\"{s3_bucket}{clip.s3_id}.mp3\"\n",
    "        # os.system(f\"aws s3 cp {clip_s3_filepath} {s3_mp3_path}\")\n",
    "\n",
    "        # add to data row list\n",
    "        assets.append(\n",
    "            {\n",
    "                \"row_data\": s3_html_url,\n",
    "                \"global_key\": request_id,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        metadata[request_id] = {\n",
    "            \"clip_id\": clip.s3_id,\n",
    "            \"clip_url\": f\"{s3_bucket_url}{clip.s3_id}.mp3\",\n",
    "            \"orig_pref\": orig_pref,\n",
    "            \"genre\": clip.genre,\n",
    "            \"tags\": clip.tags_no_gender,\n",
    "            \"lyrics\": clip.lyrics,\n",
    "            \"model\": clip.model,\n",
    "        }\n",
    "\n",
    "        global_keys.append(request_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_s3_files(html_local_paths, html_upload_paths)\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "    local_fp = [os.path.join(tmp_dir, f\"{id}.mp3\") for id in range(len(s3_paths))]\n",
    "    download_s3_files(s3_paths, local_fp)\n",
    "    upload_s3_files(local_fp, s3_upload_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save json metadata\n",
    "metadata_filepath = f\"./outputs/metadata-{test_label}.json\"\n",
    "with open(metadata_filepath, \"w\") as fp:\n",
    "    json.dump(metadata, fp)\n",
    "\n",
    "# Bulk add data rows to the dataset\n",
    "task = dataset.create_data_rows(assets)\n",
    "task.wait_till_done()\n",
    "print(task.errors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ontology_builder = lb.OntologyBuilder(\n",
    "    classifications=[\n",
    "        lb.Classification(\n",
    "            class_type=lb.Classification.Type.RADIO,\n",
    "            name=\"style\",\n",
    "            instructions=\"What gender is the singer?\",\n",
    "            options=[\n",
    "                lb.Option(value=\"Male\"),\n",
    "                lb.Option(value=\"Female\"),\n",
    "                lb.Option(value=\"No singer present\"),\n",
    "            ],\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "ontology = client.create_ontology(\n",
    "    \"Ontology HTML Annotations\", ontology_builder.asdict(), media_type=lb.MediaType.Html\n",
    ")\n",
    "\n",
    "project = client.create_project(name=\"gender-prompt-adherence3\", media_type=lb.MediaType.Html)\n",
    "\n",
    "# Setup your ontology\n",
    "project.connect_ontology(ontology)\n",
    "\n",
    "# send rows to project\n",
    "batch = project.create_batch(\n",
    "    \"first-batch-html-demo\",  # Each batch in a project must have a unique name\n",
    "    global_keys=global_keys,  # Paginated collection of data row objects, list of data row ids or global keys\n",
    "    priority=5,  # priority between 1(highest) - 5(lowest)\n",
    ")\n",
    "\n",
    "print(\"Batch: \", batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}
