{
 "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",
    "\n",
    "import tempfile\n",
    "from suno_utils.utils.s3 import upload_s3_files, download_s3_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "####### UPDATE THESE ##########\n",
    "\n",
    "# model_A = \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-auk-t1_2025_04_27-20_57_03.json\"\n",
    "# model_B = \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v4-h-s-32_2025_04_27-21_13_02.json\"\n",
    "# model_B = \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v3p5-engine-s-8_2025_04_27-21_29_57.json\"\n",
    "\n",
    "model_A = \"/home/sara/glockenspiel/suno_utils/task_eval/genius_sources_mediocre.json\"\n",
    "# model_B = \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-auk-t1_2025_04_28-20_39_36.json\"\n",
    "model_B = \"/home/sara/glockenspiel/suno_utils/task_eval/modal_runs/genre_mappings_chirp-v4-h-s-32_2025_04_29-01_47_19.json\"\n",
    "\n",
    "model_A_name = \"train\"\n",
    "model_B_name = \"v4_prod\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_gen_json(filepath):\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",
    "            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": [
    "A_df = load_gen_json(model_A)\n",
    "B_df = load_gen_json(model_B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(A_df), len(B_df))\n",
    "B_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_clip_details(clip_item):\n",
    "    return clip_item.s3_id, clip_item.lyrics, clip_item.tags, clip_item.genre\n",
    "\n",
    "\n",
    "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",
    "# is_english(A_df)\n",
    "# is_english(B_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "timestamp = datetime.now().strftime(\"%Y%m%d\")\n",
    "test_label = f\"{model_A_name}-vs-{model_B_name}-{timestamp}\"\n",
    "\n",
    "s3_bucket = f\"s3://suno-annotation-public/preference-{test_label}/\"  # this must have trailing slash\n",
    "s3_bucket_url = (\n",
    "    f\"https://suno-annotation-public.s3.amazonaws.com/preference-{test_label}\"\n",
    ")\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/ab_base.html\"\n",
    "max_num_examples = 500\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 = []  # list of assets to add to the dataset\n",
    "\n",
    "local_paths_to_write = []\n",
    "s3_paths_to_write = []\n",
    "html_local = []\n",
    "html_to_write = []\n",
    "for i in range(0, len(A_df)):\n",
    "    # get the first and second clip\n",
    "    clip_0 = A_df.iloc[i]\n",
    "    clip_1 = B_df.iloc[i]\n",
    "\n",
    "    # use request id as the unique identifier\n",
    "    request_id = f\"{clip_1.s3_id}_{test_label}\"\n",
    "\n",
    "    # get clip details\n",
    "    clip_0_id, clip_0_prompt, clip_0_tags, clip_0_genre = get_clip_details(clip_0)\n",
    "    clip_1_id, clip_1_prompt, clip_1_tags, clip_1_genre = get_clip_details(clip_1)\n",
    "    assert clip_0_prompt == clip_1_prompt\n",
    "    assert clip_0_tags == clip_1_tags\n",
    "    assert clip_0_genre == clip_1_genre\n",
    "\n",
    "    orig_pref = \"a\"  # by default the first clip is the preferred one\n",
    "\n",
    "    # randomly swap the clips\n",
    "    if np.random.rand() > 0.5:\n",
    "        # swap the clips\n",
    "        tmp_clip_id = clip_0_id\n",
    "        clip_0_id = clip_1_id\n",
    "        clip_1_id = tmp_clip_id\n",
    "        orig_pref = \"b\"\n",
    "\n",
    "    if clip_0_id.endswith(\".webm\") or clip_0_id.endswith(\".mp3\"):\n",
    "        clip_0_s3_filepath = clip_0_id\n",
    "        clip_0_id = os.path.splitext(os.path.basename(clip_0_id))[0]\n",
    "    else:\n",
    "        clip_0_s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{clip_0_id}.mp3\"\n",
    "\n",
    "    if clip_1_id.endswith(\".webm\") or clip_1_id.endswith(\".mp3\"):\n",
    "        clip_1_s3_filepath = clip_1_id\n",
    "        clip_1_id = os.path.splitext(os.path.basename(clip_1_id))[0]\n",
    "    else:\n",
    "        clip_1_s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{clip_1_id}.mp3\"\n",
    "\n",
    "    # check if the audio files exist\n",
    "    assert check_s3_file_exists(clip_0_s3_filepath)\n",
    "    assert check_s3_file_exists(clip_1_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_0_id}.mp3\"\n",
    "    soup.find(id=\"audio2\")[\"src\"] = f\"{s3_bucket_url}/{clip_1_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_0_tags}\"\n",
    "    lyrics_div.string = f\"Lyrics: {clip_0_prompt}\"\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_html_filepath = f\"{s3_bucket}{request_id}.html\"\n",
    "    s3_html_url = f\"{s3_bucket_url}/{request_id}.html\"\n",
    "    html_local.append(output_filepath)\n",
    "    html_to_write.append(s3_html_filepath)\n",
    "\n",
    "    # os.system(f\"aws s3 cp {output_filepath} {s3_html_filepath}\")\n",
    "\n",
    "    # push mp3 files to s3\n",
    "    s3_mp3_path_0 = f\"{s3_bucket}{clip_0_id}.mp3\"\n",
    "    s3_mp3_path_1 = f\"{s3_bucket}{clip_1_id}.mp3\"\n",
    "\n",
    "    local_paths_to_write.append(clip_0_s3_filepath)\n",
    "    s3_paths_to_write.append(s3_mp3_path_0)\n",
    "    local_paths_to_write.append(clip_1_s3_filepath)\n",
    "    s3_paths_to_write.append(s3_mp3_path_1)\n",
    "\n",
    "    # os.system(f\"aws s3 cp {clip_0_s3_filepath} {s3_mp3_path_0}\")\n",
    "    # os.system(f\"aws s3 cp {clip_1_s3_filepath} {s3_mp3_path_1}\")\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_a_id\": clip_0_id,\n",
    "        \"clip_b_id\": clip_1_id,\n",
    "        \"clip_a_url\": f\"{s3_bucket_url}{clip_0_id}.mp3\",\n",
    "        \"clip_b_url\": f\"{s3_bucket_url}{clip_1_id}.mp3\",\n",
    "        \"orig_pref\": orig_pref,\n",
    "        \"genre\": clip_0_genre,\n",
    "        \"tags\": clip_0_tags,\n",
    "        \"lyrics\": clip_0_prompt,\n",
    "    }\n",
    "\n",
    "    global_keys.append(request_id)\n",
    "\n",
    "    if len(assets) >= max_num_examples:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_s3_files(html_local, html_to_write, extra_args={\"ContentType\": \"text/html\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with tempfile.TemporaryDirectory() as tmp_dir:\n",
    "    local_fp = [\n",
    "        os.path.join(tmp_dir, f\"{id}.mp3\") for id in range(len(local_paths_to_write))\n",
    "    ]\n",
    "    print(download_s3_files(local_paths_to_write, local_fp))\n",
    "    print(upload_s3_files(local_fp, s3_paths_to_write))"
   ]
  },
  {
   "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(\n",
    "    name=f\"{model_A_name}-vs-{model_B_name}-preference-test\"\n",
    ")"
   ]
  },
  {
   "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=\"Which track better matches the style tag?\",\n",
    "            options=[\n",
    "                lb.Option(value=\"A\"),\n",
    "                lb.Option(value=\"B\"),\n",
    "                lb.Option(value=\"Equal\"),\n",
    "                lb.Option(value=\"I don't understand any of the style prompt\"),\n",
    "            ],\n",
    "        ),\n",
    "        lb.Classification(\n",
    "            class_type=lb.Classification.Type.RADIO,\n",
    "            name=\"lyrics\",\n",
    "            instructions=\"Which track follows the lyrics more closely?\",\n",
    "            options=[\n",
    "                lb.Option(value=\"A\"),\n",
    "                lb.Option(value=\"B\"),\n",
    "                lb.Option(value=\"Equal\"),\n",
    "            ],\n",
    "        ),\n",
    "        lb.Classification(\n",
    "            class_type=lb.Classification.Type.RADIO,\n",
    "            name=\"audio_quality\",\n",
    "            instructions=\"Which track has better audio quality?\",\n",
    "            options=[lb.Option(value=\"A\"), lb.Option(value=\"B\")],\n",
    "        ),\n",
    "        lb.Classification(\n",
    "            class_type=lb.Classification.Type.RADIO,\n",
    "            name=\"musicality\",\n",
    "            instructions=\"Which track do you prefer musically?\",\n",
    "            options=[lb.Option(value=\"A\"), lb.Option(value=\"B\")],\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(\n",
    "    name=f\"{model_A_name}-vs-{model_B_name}-preference\", media_type=lb.MediaType.Html\n",
    ")\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
}
