{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "import time\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from suno_utils.tasks.ear import load_model\n",
    "# score semantc\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = load_model(\"/app/suno/christian/checkpoints/ear-v2/2025-04-04_09-51-12_s9595/last_ckpt.pt\", compile=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_npz_from_s3(gen_id ):\n",
    "    s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "    data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "    if \"v3.0_raw\" in data:\n",
    "        codes = data[\"v3.0_raw\"]\n",
    "    elif \"v3.5_raw\" in data:\n",
    "        codes = data[\"v3.5_raw\"]\n",
    "    elif \"v4.0_raw\" in data:\n",
    "        codes = data[\"v4.0_raw\"]\n",
    "    elif \"v5.0_raw\" in data:\n",
    "        codes = data[\"v5.0_raw\"]\n",
    "    else:\n",
    "        raise ValueError(\"No codes found\")  \n",
    "    \n",
    "    return codes[:, 0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load tony's data from auk\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "filepath = \"/home/tony/Data/Preference/auk/interesting_clips_exp_20250404_auk_test.pkl\"\n",
    "df = pd.read_pickle(filepath)\n",
    "\n",
    "auk_id = \"chirp-v4-6b-t-04\"\n",
    "v4_id = \"chirp-v4-h-s-32\"\n",
    "\n",
    "model_scores = {\n",
    "    \"chirp-v4-6b-t-04\": [],\n",
    "    \"chirp-v4-h-s-32\": [],\n",
    "    \"chirp-v4-h-t-6\" : [],\n",
    "}\n",
    "\n",
    "# iterate over pairs of rows, odd index\n",
    "row_idx = np.arange(len(df))\n",
    "pbar = tqdm(row_idx)\n",
    "for i in pbar:\n",
    "    row_1 = df.iloc[i]\n",
    "    row_2 = df.iloc[i+1]\n",
    "\n",
    "    # get the npz file from s3\n",
    "    row1_codes = load_npz_from_s3(row_1[\"id\"])\n",
    "    row2_codes = load_npz_from_s3(row_2[\"id\"])\n",
    "\n",
    "    # score them \n",
    "    \n",
    "    row1_score = model.get_score_batch_semantic(torch.from_numpy(row1_codes).long().unsqueeze(0).cuda())\n",
    "    row2_score = model.get_score_batch_semantic(torch.from_numpy(row2_codes).long().unsqueeze(0).cuda())\n",
    "\n",
    "    #print(f\"{row_1['model_name']} {row_1['preference']} {row1_score}\")\n",
    "    #print(f\"{row_2['model_name']} {row_2['preference']} {row2_score}\")\n",
    "\n",
    "    model_scores[row_1[\"model_name\"]].append(row1_score.item())\n",
    "    model_scores[row_2[\"model_name\"]].append(row2_score.item())\n",
    "\n",
    "    if i > 100:\n",
    "        break\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert tensors to items\n",
    "model_scores = {k: [v.item() for v in v] for k, v in model_scores.items()}\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# plot histogram of model scores\n",
    "bins = np.linspace(0, 15, 25)\n",
    "plt.hist(model_scores[\"chirp-v4-6b-t-04\"], bins=bins, label=\"chirp-v4-6b-t-04\")\n",
    "plt.hist(model_scores[\"chirp-v4-h-s-32\"], bins=bins, label=\"chirp-v4-h-s-32\")\n",
    "plt.hist(model_scores[\"chirp-v4-h-t-6\"], bins=bins, label=\"chirp-v4-h-t-6\")\n",
    "plt.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# score semantc\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "gen_id = \"32784d29-ff57-4769-a9b9-b853d8af7953\"\n",
    "\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "if \"v3.0_raw\" in data:\n",
    "    codes = data[\"v3.0_raw\"]\n",
    "elif \"v3.5_raw\" in data:\n",
    "    codes = data[\"v3.5_raw\"]\n",
    "elif \"v4.0_raw\" in data:\n",
    "    codes = data[\"v4.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n",
    "\n",
    "gen_semantic_codes = torch.from_numpy(codes[:, 0]).long()#.cuda()\n",
    "print(gen_semantic_codes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_score = model.get_score_batch_semantic(gen_semantic_codes.unsqueeze(0).cuda())\n",
    "print(gen_score)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}