{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from suno_utils.utils.text import write_jsonl, read_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filepath = \"/home/tony/Data/Preference/auk/interesting_clips_exp_20250422_auk_t1.pkl\"\n",
    "df = pd.read_pickle(filepath)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get only positive clips\n",
    "pos_df = df[df[\"preference\"] == True]\n",
    "print(len(pos_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse the metadata for clips and create a new dataframe\n",
    "import json\n",
    "\n",
    "# Create a new dataframe with just the columns we need\n",
    "new_df = pd.DataFrame(columns=['id', 'tags', 'text'])\n",
    "\n",
    "# For each row, parse the metadata and extract tags and text\n",
    "rows_to_add = []\n",
    "for idx, row in pos_df.iterrows():\n",
    "    try:\n",
    "        metadata = row[\"metadata\"]\n",
    "        if isinstance(metadata, str):\n",
    "            metadata = json.loads(metadata)\n",
    "        \n",
    "        # Extract data\n",
    "        tags = metadata.get(\"tags\", None)\n",
    "        text = metadata.get(\"prompt\", metadata.get(\"text\", None))\n",
    "        \n",
    "        # Add to list of rows\n",
    "        rows_to_add.append({\n",
    "            'id': row['id'],\n",
    "            'tags': tags,\n",
    "            'text': text\n",
    "        })\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing row {idx}: {e}\")\n",
    "\n",
    "# Create dataframe from collected rows\n",
    "new_df = pd.concat([new_df, pd.DataFrame(rows_to_add)], ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save new df to jsonl\n",
    "new_metas = []\n",
    "for idx, row in new_df.iterrows():\n",
    "    new_metas.append({\n",
    "        \"id\": row[\"id\"],\n",
    "        \"text\": row[\"text\"],\n",
    "        \"tags\": row[\"tags\"]\n",
    "    })\n",
    "\n",
    "\n",
    "write_jsonl(new_metas, \"/home/christian/code/christian/metadata/sft/auk_clips_up_u_1_20241201_pos.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meta_id = new_df.iloc[0][\"id\"]\n",
    "text = new_df.iloc[0][\"text\"]\n",
    "tags = new_df.iloc[0][\"tags\"]\n",
    "print(meta_id)\n",
    "print(text)\n",
    "print(tags)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}