{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import ast\n",
    "import glob\n",
    "import boto3\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import labelbox as lb\n",
    "\n",
    "\n",
    "from tqdm import tqdm\n",
    "from langdetect import detect\n",
    "from bs4 import BeautifulSoup\n",
    "#from better_profanity import profanity\n",
    "from suno_utils.utils.s3 import check_s3_file_exists, read_from_s3\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_s3_files_recursive(bucket_name, prefix, max_keys: int = 100000):\n",
    "    all_files = []\n",
    "    continuation_token = None\n",
    "    print(bucket_name, prefix)\n",
    "\n",
    "    while True:\n",
    "        # Prepare the arguments for the request\n",
    "        list_kwargs = {\n",
    "            \"Bucket\": bucket_name,\n",
    "            \"Prefix\": prefix,  # List objects under this prefix, or leave blank for all objects\n",
    "        }\n",
    "\n",
    "        if continuation_token:\n",
    "            list_kwargs[\"ContinuationToken\"] = continuation_token\n",
    "\n",
    "        # Make the request to list objects\n",
    "        response = s3.list_objects_v2(**list_kwargs)\n",
    "\n",
    "        # Collect the file keys\n",
    "        all_files += [obj[\"Key\"] for obj in response.get(\"Contents\", [])]\n",
    "\n",
    "        # Check if more results are available\n",
    "        if response.get(\"IsTruncated\"):  # True if there are more results to fetch\n",
    "            continuation_token = response[\"NextContinuationToken\"]\n",
    "        else:\n",
    "            break  # No more results to fetch\n",
    "\n",
    "    return all_files\n",
    "\n",
    "def get_s3_files(bucket_name, prefix, max_keys: int = 100000):\n",
    "    \"\"\"\n",
    "    List S3 files in the specified prefix level only (non-recursive).\n",
    "    \n",
    "    Args:\n",
    "        bucket_name (str): Name of the S3 bucket\n",
    "        prefix (str): Prefix to list objects under\n",
    "        max_keys (int): Maximum number of keys to return\n",
    "        \n",
    "    Returns:\n",
    "        list: List of file keys at the specified prefix level\n",
    "    \"\"\"\n",
    "    all_files = []\n",
    "    continuation_token = None\n",
    "    \n",
    "    # Normalize prefix to not end with / unless it's empty\n",
    "    if prefix and prefix != '/':\n",
    "        prefix = prefix.rstrip('/')\n",
    "    \n",
    "    while len(all_files) < max_keys:\n",
    "        # Prepare the arguments for the request\n",
    "        list_kwargs = {\n",
    "            \"Bucket\": bucket_name,\n",
    "            \"Prefix\": prefix,\n",
    "            \"Delimiter\": '/',\n",
    "            \"MaxKeys\": min(1000, max_keys - len(all_files))  # S3 max is 1000 per request\n",
    "        }\n",
    "        \n",
    "        if continuation_token:\n",
    "            list_kwargs[\"ContinuationToken\"] = continuation_token\n",
    "            \n",
    "        # Make the request to list objects\n",
    "        response = s3.list_objects_v2(**list_kwargs)\n",
    "        \n",
    "        # Process the direct contents (files)\n",
    "        for obj in response.get(\"Contents\", []):\n",
    "            key = obj[\"Key\"]\n",
    "            \n",
    "            # Skip if this is the prefix directory itself\n",
    "            if key == prefix or key == prefix + '/':\n",
    "                continue\n",
    "                \n",
    "            # For files directly in this level\n",
    "            all_files.append(key)\n",
    "            \n",
    "        # Check if more results are available\n",
    "        if not response.get(\"IsTruncated\"):\n",
    "            break\n",
    "            \n",
    "        continuation_token = response.get(\"NextContinuationToken\")\n",
    "        \n",
    "    return all_files[:max_keys]\n",
    "\n",
    "def list_s3_directories(bucket_name: str, prefix: str = ''):\n",
    "    \"\"\"\n",
    "    List all directories (prefixes) in an S3 bucket.\n",
    "    \n",
    "    Args:\n",
    "        bucket_name (str): Name of the S3 bucket\n",
    "        prefix (str): Optional prefix to filter results (like a directory path)\n",
    "        \n",
    "    Returns:\n",
    "        List[str]: List of directory paths (prefixes)\n",
    "    \"\"\"\n",
    "    s3_client = boto3.client('s3')\n",
    "    directories = set()\n",
    "    \n",
    "    # Use paginator to handle buckets with many objects\n",
    "    paginator = s3_client.get_paginator('list_objects_v2')\n",
    "    page_iterator = paginator.paginate(\n",
    "        Bucket=bucket_name,\n",
    "        Prefix=prefix,\n",
    "        Delimiter='/'\n",
    "    )\n",
    "    \n",
    "    # Collect all prefixes (directories)\n",
    "    for page in page_iterator:\n",
    "        # Get common prefixes (directories)\n",
    "        if 'CommonPrefixes' in page:\n",
    "            for prefix_obj in page['CommonPrefixes']:\n",
    "                directories.add(prefix_obj['Prefix'])\n",
    "                \n",
    "        # Also check Contents for any directory-like objects\n",
    "        if 'Contents' in page:\n",
    "            for obj in page['Contents']:\n",
    "                key = obj['Key']\n",
    "                # If the key contains a slash, add the directory part\n",
    "                if '/' in key:\n",
    "                    directory = key.rsplit('/', 1)[0] + '/'\n",
    "                    directories.add(directory)\n",
    "    \n",
    "    return sorted(list(directories))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_name = \"suno-data\"\n",
    "# base_dir = \"christian/data/upsample_100z_v1\"\n",
    "# output_name = \"v2\"\n",
    "base_dir = \"christian/data/upsample_v4_t_5_20241018\"\n",
    "output_name = \"25hz_20241031_v1/\"\n",
    "\n",
    "base_metas_path = os.path.join(\"s3://\", bucket_name, base_dir, \"metas.jsonl\")\n",
    "base_metas = read_from_s3(base_metas_path, read_f=read_jsonl)\n",
    "print(len(base_metas))\n",
    "\n",
    "# find all files on s3 with the pattern\n",
    "s3 = boto3.client(\"s3\")\n",
    "dir_paths = list_s3_directories(bucket_name, f\"{base_dir}/{output_name}\")\n",
    "print(\"total files on s3: \", len(dir_paths))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_metas_map = {meta[\"id\"]: meta for meta in base_metas}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "pairs = []\n",
    "# collect examples \n",
    "for dir_path in tqdm(dir_paths):\n",
    "    meta = base_metas_map[dir_path.strip(\"/\").split(\"/\")[-1]]\n",
    "\n",
    "    # load the pair_info.json\n",
    "    pair_info_path = f\"s3://{bucket_name}/{dir_path}pair_info.json\"\n",
    "    pair_info = read_from_s3(pair_info_path)\n",
    "    pair_info = json.loads(pair_info)\n",
    "\n",
    "    # detect language\n",
    "    lyrics = meta.get(\"text\", None)\n",
    "    if lyrics is not None and lyrics != \"\":\n",
    "        lang = detect(lyrics)\n",
    "        print(f\"language: {lang}\")\n",
    "        if lang != \"en\":\n",
    "            print(f\"skipping {meta['id']} because language is not english\")\n",
    "            continue\n",
    "\n",
    "    # get mp3 paths of the two examples\n",
    "    a_filename = pair_info[\"a_filename\"]\n",
    "    b_filename = pair_info[\"b_filename\"]\n",
    "\n",
    "    a_s3_path = f\"s3://{bucket_name}/{dir_path}{a_filename}.mp3\"\n",
    "    b_s3_path = f\"s3://{bucket_name}/{dir_path}{b_filename}.mp3\"\n",
    "\n",
    "    pairs.append({\"id\": meta[\"id\"], \"a_s3_path\": a_s3_path, \"b_s3_path\": b_s3_path, \"a_filename\": a_filename, \"b_filename\": b_filename, \"a_seed\": pair_info[\"a_seed\"], \"b_seed\": pair_info[\"b_seed\"], \"lyrics\": meta.get(\"text_aligned\", \"\"), \"tags\": meta.get(\"tags\", \"\")})\n",
    "\n",
    "    if len(pairs) >= 100:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a dataset\n",
    "client = lb.Client(api_key=\"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VySWQiOiJjbHppbmJnY2wwMDYyMDd5bWg2enhiMTd6Iiwib3JnYW5pemF0aW9uSWQiOiJjbHppbmJnY2QwMDYxMDd5bWM4cjM5cHdzIiwiYXBpS2V5SWQiOiJjbHp3cmU3ZnQwYjFzMDd6aWRrNTFnb21mIiwic2VjcmV0IjoiMGU5M2MwNmU4ZWI1Y2Y3NTlmNTk5YTk5MzIwOTU5MzQiLCJpYXQiOjE3MjM4MTU3NTcsImV4cCI6MjM1NDk2Nzc1N30.Z6gZwlzQ85KrqGOydqCdo1RVmUhOEntpev2HhNso4PU\")\n",
    "dataset = client.create_dataset(name=\"upsample-compare-ab-preference-test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_clip_details(clip_item):\n",
    "    s3_id = clip_item[\"s3_id\"]\n",
    "    string_data = clip_item[\"metadata\"]\n",
    "    dict_data = ast.literal_eval(string_data)\n",
    "    lyrics = dict_data[\"prompt\"]\n",
    "    tags = dict_data[\"tags\"]\n",
    "\n",
    "    return s3_id, lyrics, tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "timestamp = 20241104\n",
    "# s3://suno-annotation-public/upsample-ab-preference-20241022/\n",
    "s3_bucket = f\"s3://suno-annotation-public/upsample-compare-ab-preference-{timestamp}/\" # this must have trailing slash\n",
    "s3_bucket_url = f\"https://suno-annotation-public.s3.amazonaws.com/upsample-compare-ab-preference-{timestamp}\"\n",
    "output_html_dir = f\"/home/christian/code/christian/outputs/html-upsample-compare-ab-preference-{timestamp}\"\n",
    "os.makedirs(output_html_dir, exist_ok=True)\n",
    "\n",
    "use_prompt = False\n",
    "if use_prompt:\n",
    "    base_html_filepath = \"/home/christian/code/christian/labelbox/templates/ab_base.html\"\n",
    "else:\n",
    "    base_html_filepath = \"/home/christian/code/christian/labelbox/templates/ab_base_no_prompt.html\"\n",
    "\n",
    "max_num_examples = 100\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",
    "for pair in tqdm(pairs):\n",
    "    request_id = pair[\"id\"]\n",
    "    lyrics = pair[\"lyrics\"]\n",
    "    tags = pair[\"tags\"]\n",
    "    \n",
    "    clip_a_s3_filepath = f\"{pair['a_s3_path']}\"\n",
    "    clip_b_s3_filepath = f\"{pair['b_s3_path']}\"\n",
    "    clip_a_basename = pair[\"a_filename\"].split(\"/\")[-1]\n",
    "    clip_b_basename = pair[\"b_filename\"].split(\"/\")[-1]\n",
    "\n",
    "    # copy files to local temporary directory\n",
    "    temp_dir = \"/home/christian/code/christian/notebooks/outputs/temp\"\n",
    "    os.makedirs(temp_dir, exist_ok=True)\n",
    "    os.system(f\"aws s3 cp {clip_a_s3_filepath} {temp_dir}/{clip_a_basename}.mp3\")\n",
    "    os.system(f\"aws s3 cp {clip_b_s3_filepath} {temp_dir}/{clip_b_basename}.mp3\")\n",
    "\n",
    "    # then loudness normalize to -16db LUFS with ffmpeg force overwrite with 320 kbps bitrate\n",
    "    os.system(f\"ffmpeg -i {temp_dir}/{clip_a_basename}.mp3 -filter:a loudnorm=I=-16:TP=-1.5:LRA=11 -c:a libmp3lame {temp_dir}/{clip_a_basename}_norm.mp3 -y -b:a 320k\")\n",
    "    os.system(f\"ffmpeg -i {temp_dir}/{clip_b_basename}.mp3 -filter:a loudnorm=I=-16:TP=-1.5:LRA=11 -c:a libmp3lame {temp_dir}/{clip_b_basename}_norm.mp3 -y -b:a 320k\")\n",
    "    # copy audio files to s3 bucket\n",
    "    print(f\"copying {temp_dir}/{clip_a_basename}_norm.mp3 to {s3_bucket}{clip_a_basename}_norm.mp3\")\n",
    "    print(f\"copying {temp_dir}/{clip_b_basename}_norm.mp3 to {s3_bucket}{clip_b_basename}_norm.mp3\")\n",
    "    os.system(f\"aws s3 cp {temp_dir}/{clip_a_basename}_norm.mp3 {s3_bucket}{clip_a_basename}_norm.mp3\")\n",
    "    os.system(f\"aws s3 cp {temp_dir}/{clip_b_basename}_norm.mp3 {s3_bucket}{clip_b_basename}_norm.mp3\")\n",
    "\n",
    "    # edit the html file to add audios and prompt information\n",
    "    soup.find(id=\"audio1\")['src'] = f\"{s3_bucket_url}/{clip_a_basename}_norm.mp3\"\n",
    "    soup.find(id=\"audio2\")['src'] = f\"{s3_bucket_url}/{clip_b_basename}_norm.mp3\"\n",
    "\n",
    "    # Find the elements by ID\n",
    "    if use_prompt:\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: {tags}\"\n",
    "        lyrics_div.string = f\"Lyrics: {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_html_filepath = f\"{s3_bucket}{request_id}.html\"\n",
    "    s3_html_url = f\"{s3_bucket_url}/{request_id}.html\"\n",
    "    print(s3_html_url)\n",
    "    os.system(f\"aws s3 cp {output_filepath} {s3_html_filepath}\")\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\": pair[\"a_seed\"],\n",
    "        \"clip_b_id\": pair[\"b_seed\"],\n",
    "        \"clip_a_url\" : f\"{s3_bucket_url}/{pair['a_filename']}.mp3\",\n",
    "        \"clip_b_url\" : f\"{s3_bucket_url}/{pair['b_filename']}.mp3\",\n",
    "    }\n",
    "\n",
    "    global_keys.append(request_id)\n",
    "    \n",
    "    if len(assets) >= max_num_examples:\n",
    "        break\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save json metadata\n",
    "metadata_filepath = f\"/home/christian/code/christian/labelbox/metadata/metadata-upsample-compare-ab-preference-{timestamp}.json\"\n",
    "with open(metadata_filepath, \"w\") as fp:\n",
    "    json.dump(metadata, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": [
    "print(metadata[\"026eeb93-dbd8-4967-bffe-54547b828bac\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup Ontology and Create Labeling Project"
   ]
  },
  {
   "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=\"Preference\", \n",
    "        options=[\n",
    "            lb.Option(value=\"A\"),\n",
    "            lb.Option(value=\"B\"),\n",
    "        ]\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "ontology = client.create_ontology(\"Ontology HTML Annotations\", ontology_builder.asdict(), media_type=lb.MediaType.Html)\n",
    "\n",
    "project = client.create_project(name=\"music-ab-preference\", \n",
    "                                    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": "markdown",
   "metadata": {},
   "source": [
    "# Covers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# find all examples \n",
    "local_dir = \"/home/christian/code/christian/notebooks/outputs/covers-20240816\"\n",
    "s3_bucket = \"s3://suno-annotation-public/covers-20240816\"\n",
    "s3_bucket_url = \"https://suno-annotation-public.s3.amazonaws.com/covers-20240816\"\n",
    "\n",
    "# find all json files \n",
    "metas = glob.glob(os.path.join(local_dir, \"*.json\"))\n",
    "print(f\"Found {len(metas)} examples...\")\n",
    "\n",
    "metas = metas[:10] # just the first 10 for now\n",
    "\n",
    "assets = [] # list of assets to add to the dataset\n",
    "\n",
    "for meta in metas:\n",
    "    with open(meta, \"r\") as f:\n",
    "        data = json.load(f)\n",
    "\n",
    "    meta_id = data[\"uid\"]\n",
    "\n",
    "    # upload the audio files to s3 bucket\n",
    "    #os.system(f\"aws s3 cp {local_dir}/{meta_id}-source.mp3 {s3_bucket}/\")\n",
    "    #os.system(f\"aws s3 cp {local_dir}/{meta_id}-cover-0.mp3 {s3_bucket}/\")\n",
    "    #os.system(f\"aws s3 cp {local_dir}/{meta_id}-cover-1.mp3 {s3_bucket}/\")\n",
    "\n",
    "    # build paths to audio files with this uid\n",
    "    s3_src_a_path = os.path.join(s3_bucket_url, f\"{meta_id}-source.mp3\")\n",
    "    s3_gen_a_path = os.path.join(s3_bucket_url, f\"{meta_id}-cover-0.mp3\")\n",
    "    s3_gen_b_path = os.path.join(s3_bucket_url, f\"{meta_id}-cover-1.mp3\")\n",
    "\n",
    "    target_tags = data[\"target_tags\"]\n",
    "    print(\"target tags:\", target_tags)\n",
    "    lyrics = data[\"lyrics\"]\n",
    "\n",
    "    assets.append(\n",
    "        {\n",
    "            \"global_key\": meta_id,\n",
    "            \"row_data\": s3_html_url,\n",
    "        \n",
    "        }\n",
    "    )\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": [
    "# Create data payload\n",
    "# Use global key, a unique ID to identify an asset throughout Labelbox workflow. Learn more: https://docs.labelbox.com/docs/global-keys\n",
    "# You can add metadata fields to your data rows. Learn more: https://docs.labelbox.com/docs/import-metadata\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env2",
   "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
}
