{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# to start we can probably just random sample some clips discogs_subset \n",
    "# we won't use alignments for the first pass\n",
    "import os\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl\n",
    "\n",
    "metas_dir = \"/app/suno/tmp\"\n",
    "dicsogs_subset_metas = read_jsonl(os.path.join(metas_dir, \"raw_discogs_subset_metas.jsonl\"))\n",
    "print(len(dicsogs_subset_metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "# now we can apply some basic filter and sample a subset \n",
    "# we can use the sampled_tracks subset with ids\n",
    "sampled_tracks_filepath = \"/home/christian/code/christian/metadata/sampled_tracks.json\"\n",
    "with open(sampled_tracks_filepath, \"r\") as f:\n",
    "    sampled_tracks = json.load(f)\n",
    "\n",
    "discogs_sampled_ids = set(sampled_tracks[\"discogs_subset\"])\n",
    "print(len(discogs_sampled_ids))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# now we can sample a subset of these \n",
    "discogs_subset_sampled_metas = [meta for meta in dicsogs_subset_metas if meta[\"id\"] in discogs_sampled_ids]\n",
    "num_sampled = len(discogs_subset_sampled_metas)\n",
    "print(f\"num sampled: {num_sampled} out of {len(dicsogs_subset_metas)} ({num_sampled/len(dicsogs_subset_metas)*100:.2f}%)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "discogs_subset_sampled_metas[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "# now create new jsonl file with these \n",
    "output_file = os.path.join(\"/home/christian/code/christian/metadata\", \"discogs_subset_sampled_metas.jsonl\")\n",
    "\n",
    "# but we will want to keep only the relevant fields\n",
    "new_metas = []\n",
    "for meta in tqdm(discogs_subset_sampled_metas):\n",
    "\n",
    "    sos_genre = meta.get(\"sos_genre\", [])\n",
    "    rym_genres = meta.get(\"rym_genres\", [])\n",
    "    duration_s = meta.get(\"duration_s\", None)\n",
    "    year = meta.get(\"year\", None)\n",
    "    musicbrainz_metadata = meta.get(\"musicbrainz_album\", None)\n",
    "    # compile into list of tags\n",
    "    tags = []\n",
    "    tags.extend(sos_genre)\n",
    "\n",
    "    if len(rym_genres) > 0:\n",
    "        rym_tag_list = rym_genres.split(\",\")\n",
    "        rym_tag_list = [tag.strip() for tag in rym_tag_list]\n",
    "        tags.extend(rym_tag_list)\n",
    "\n",
    "    if year is not None:\n",
    "        tags.append(f\"{year}\")\n",
    "\n",
    "    if musicbrainz_metadata is not None:\n",
    "        recording_genres = musicbrainz_metadata.get(\"recording_genres\", [])\n",
    "        recording_tags = musicbrainz_metadata.get(\"recording_tags\", [])\n",
    "        tags.extend(recording_genres)\n",
    "        tags.extend(recording_tags)\n",
    "\n",
    "    new_meta = {\n",
    "        \"id\": meta[\"id\"],\n",
    "        \"title\": meta[\"title\"],\n",
    "        \"artists\": meta[\"artists\"],\n",
    "        \"s3_filepath\": meta[\"s3_filepath\"],\n",
    "        \"text\": meta[\"lyrics\"],\n",
    "        \"tags\": tags,\n",
    "        \"duration_s\": meta[\"duration_s\"],\n",
    "    }\n",
    "    new_metas.append(new_meta)\n",
    "\n",
    "\n",
    "write_jsonl(new_metas, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load preference data for getting prompts and semantic\n",
    "pkl_filepath = \"/home/tony/Data/Preference/up_v1/interesting_clips_up_u_1_20241201_full.pkl\"\n",
    "df = pd.read_pickle(pkl_filepath)\n",
    "print(len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "work_items = read_from_s3(\n",
    "    \"s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl\",\n",
    "    read_f=read_jsonl,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(work_items[0][\"id\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(df.iloc[0][\"metadata\"][\"prompt\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "metas = []\n",
    "\n",
    "for idx in tqdm(range(0, len(df), 2)):\n",
    "    # get the two rows\n",
    "    row1 = df.iloc[idx] # first is negative\n",
    "    row2 = df.iloc[idx + 1] # second is positive\n",
    "    # get the text and tags\n",
    "\n",
    "    id2 = row2[\"id_x\"]\n",
    "    lyrics = row2[\"metadata\"][\"prompt\"]\n",
    "    tags = row2[\"metadata\"][\"tags\"]\n",
    "    \n",
    "    metas.append({\n",
    "        \"id\": id2,\n",
    "        \"text\": lyrics,\n",
    "        \"tags\": tags,\n",
    "    })\n",
    "\n",
    "print(len(metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save metas\n",
    "import os\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl\n",
    "\n",
    "output_dir = \"/home/christian/code/christian/metadata/sft\"\n",
    "output_file = os.path.join(output_dir, \"pos_interesting_clips_up_u_1_20241201_full.jsonl\")\n",
    "write_jsonl(metas, output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clone_diff",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}