{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#data_path = \"/home/tony/Data/Preference/30b_v5/interesting_clips_v4_t_5_20241015_full_with_cer_pos_gen.pkl\"\n",
    "data_path = \"/home/tony/Data/Preference/30b_v5/interesting_clips_v4_t_5_20241018_full_with_cer_pos_gen.pkl\"\n",
    "npz_root_dir = \"/app/suno/data/dpo/30b_npz/\"\n",
    "json_root_dir = \"/app/suno/data/dpo/30b_json/\"\n",
    "df = pd.read_pickle(data_path)\n",
    "print(f\"df shape: {df.shape}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = []\n",
    "\n",
    "for idx, (row_idx, row) in enumerate(df.iterrows()):\n",
    "\n",
    "\n",
    "    text_aligned = \"\"\n",
    "    text = \"\"\n",
    "\n",
    "    alignment_path = os.path.join(json_root_dir, f\"{row.id_x}_hoot.json\")\n",
    "    if os.path.exists(alignment_path):\n",
    "\n",
    "        # load alignment from json\n",
    "        with open(os.path.join(json_root_dir, f\"{row.id_x}_hoot.json\"), \"r\") as f:\n",
    "            alignment = json.load(f)\n",
    "\n",
    "        for elem in alignment:\n",
    "            if \"word\" in elem and \"start_s\" in elem:\n",
    "                if float(elem[\"start_s\"]) < 30.0:\n",
    "                    text_aligned += elem[\"word\"]\n",
    "                text += elem[\"word\"]\n",
    "\n",
    "        # get a string with the first 30s of lyrics in alignment\n",
    "\n",
    "    meta = {\n",
    "        \"id\": str(row.id_x),\n",
    "        \"text_aligned\": text_aligned,\n",
    "        \"text\" : text,\n",
    "        \"tags\" : row.tags,\n",
    "        \"start_s\" : 0.0,\n",
    "        \"end_s\" : 30.0  \n",
    "    }\n",
    "    \n",
    "    metas.append(meta)\n",
    "\n",
    "write_jsonl(metas, \"/home/christian/code/christian/metadata/interesting_clips_v4_t_5_20241018_full_with_cer_pos_gen.jsonl\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# aws s3 cp /home/christian/code/christian/metadata/interesting_clips_v4_t_5_20241018_full_with_cer_pos_gen.jsonl s3://suno-data/christian/data/upsample_100hz_v4_t_5_20241018/metas.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(metas))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Genius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl\")\n",
    "print(len(metas))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "random.seed(42)\n",
    "new_metas = []\n",
    "random.shuffle(metas)\n",
    "for meta in metas:\n",
    "    if \"lyrics\" in meta and len(meta[\"lyrics\"]) > 0:\n",
    "        new_metas.append(meta)\n",
    "    if len(new_metas) >= 20000:\n",
    "        break\n",
    "\n",
    "print(len(new_metas), len(metas))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "write_jsonl(new_metas, \"/home/christian/code/christian/metadata/genius_hq_filtered_20k.jsonl\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "for key, value in new_metas[0].items():\n",
    "    print(f\"{key}: {value}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}