{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import uuid\n",
    "import polars as pl\n",
    "from suno_utils.utils.text import read_jsonl, write_jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_name = \"discogs_subset\"\n",
    "\n",
    "# load relevant dataset titles\n",
    "filepath = f\"/home/christian/code/christian/metadata/dedup/{dataset_name}_titles.json\"\n",
    "with open(filepath, \"r\") as f:\n",
    "    titles = json.load(f)\n",
    "\n",
    "print(len(titles))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "client = OpenAI(api_key=\"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\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are a music expert with an expansive knowledge of music.\n",
    "You will be provided with a youtube title from a song and will create a version of the title that represents that original song \n",
    "but removes any reference to bands, artists, performers, song names, compositions, album titles, or similar identifying information. \n",
    "If there is information about the original song, track, or artist remove that, but retain important information about the content where possible.\n",
    "Only respond with the enhanced title.\n",
    "\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "\n",
    "#MAX_METAS = 400_000\n",
    "# Define the results file path\n",
    "results_file = f'/home/christian/code/christian/metadata/tagging/gpt-4_1-title-tagging_results_{dataset_name}.json'\n",
    "\n",
    "# Initialize results dictionary\n",
    "results = {}\n",
    "\n",
    "# Load existing results if file exists\n",
    "if os.path.exists(results_file):\n",
    "    with open(results_file, 'r') as f:\n",
    "        results = json.load(f)\n",
    "    print(f\"Loaded {len(results)} existing results from {results_file}\")\n",
    "\n",
    "# Filter out metas that have already been processed\n",
    "titles_to_process = [meta for meta in titles if meta[\"id\"] not in results]\n",
    "print(f\"Processing {len(titles_to_process)} new items out of {len(titles)} total\")\n",
    "\n",
    "\n",
    "# we will create a jsonl file for batch api\n",
    "batch_file = []\n",
    "\n",
    "# split the metas in to list of chunk_size\n",
    "chunk_size = 50\n",
    "chunked_titles = [titles_to_process[i:i+chunk_size] for i in range(0, len(titles_to_process), chunk_size)]\n",
    "print(len(chunked_titles))\n",
    "\n",
    "# to start lets just do 10 chunks\n",
    "\n",
    "for chunk in chunked_titles[:10]:\n",
    "    # create a uuid for each chunk\n",
    "    chunk_uuid = str(uuid.uuid4())\n",
    "    chunk_info = \"\"\n",
    "    for title in chunk:\n",
    "        chunk_info += f\"{title['id']}: {title['title']}\\n\"\n",
    "\n",
    "    batch_file.append({\n",
    "        \"custom_id\": f\"title-{chunk_uuid}\",\n",
    "        \"method\": \"POST\",\n",
    "        \"url\": \"/v1/chat/completions\",\n",
    "        \"body\": {\n",
    "            \"model\": \"gpt-4.1-mini\",\n",
    "            \"messages\": [{\"role\": \"system\", \"content\": system_prompt}, {\"role\": \"user\", \"content\": chunk_info}],\n",
    "            \"max_tokens\": 4096\n",
    "        }\n",
    "    })\n",
    "\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(batch_file[4][\"body\"][\"messages\"][1][\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(batch_file)\n",
    "batch_filepath = \"/home/christian/code/christian/metadata/tagging/batchinput_t8.jsonl\"\n",
    "# save the batch file to a jsonl file\n",
    "with open(batch_filepath, \"w\") as f:\n",
    "    for item in batch_file:\n",
    "        f.write(json.dumps(item) + \"\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload the batch file to the client\n",
    "batch_input_file = client.files.create(\n",
    "    file=open(batch_filepath, \"rb\"),\n",
    "    purpose=\"batch\"\n",
    ")\n",
    "\n",
    "print(batch_input_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a batch job\n",
    "batch_input_file_id = batch_input_file.id\n",
    "client.batches.create(\n",
    "    input_file_id=batch_input_file_id,\n",
    "    endpoint=\"/v1/chat/completions\",\n",
    "    completion_window=\"24h\",\n",
    "    metadata={\n",
    "        \"description\": \"batch gpt-4.1 audio tagging genius (10)\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = client.batches.retrieve(\"batch_681175ae91dc819086cd6d074d182d01\")\n",
    "print(batch.request_counts)\n",
    "print(batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve the batch job\n",
    "file_response = client.files.content(batch.output_file_id)\n",
    "#print(file_response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse the JSONL string into a list of dictionaries\n",
    "responses = [json.loads(line) for line in file_response.text.split('\\n') if line.strip()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for response in tqdm(responses):\n",
    "    item_id = response[\"custom_id\"].replace(\"meta-\", \"\")\n",
    "    content = response[\"response\"][\"body\"][\"choices\"][0][\"message\"][\"content\"]\n",
    "    # parse the content as json\n",
    "    content = json.loads(content)\n",
    "    # get the gpt_tag\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# file_response is a string, we need to parse it as jsonl \n",
    "import json\n",
    "from tqdm import tqdm\n",
    "\n",
    "dataset_name = \"imslp\"\n",
    "\n",
    "# Parse the JSONL string into a list of dictionaries\n",
    "responses = [json.loads(line) for line in file_response.text.split('\\n') if line.strip()]\n",
    "\n",
    "results_file = f'/home/christian/code/christian/metadata/tagging/gpt-4_1-audio-tagging_results_{dataset_name}.json'\n",
    "\n",
    "# Initialize results dictionary\n",
    "results = {}\n",
    "\n",
    "# Load existing results if file exists\n",
    "if os.path.exists(results_file):\n",
    "    with open(results_file, 'r') as f:\n",
    "        results = json.load(f)\n",
    "    print(f\"Loaded {len(results)} existing results from {results_file}\")\n",
    "\n",
    "for response in tqdm(responses):\n",
    "    item_id = response[\"custom_id\"].replace(\"meta-\", \"\")\n",
    "    content = response[\"response\"][\"body\"][\"choices\"][0][\"message\"][\"content\"]\n",
    "    \n",
    "    # Add to dictionary with item_id as key\n",
    "    results[item_id] = {\n",
    "        \"gpt_tag\": content\n",
    "    }\n",
    "    \n",
    "    # Also print for verification\n",
    "    #print(f\"ID: {item_id}\")\n",
    "    #print(f\"Content: {content}\")\n",
    "    #print()\n",
    "\n",
    "print(f\"total results: {len(results)}\")\n",
    "\n",
    "# Write the results to a normal JSON file\n",
    "with open(results_file, 'w') as f:\n",
    "    json.dump(results, f, indent=2)\n",
    "\n",
    "print(f\"Results written to {results_file}\")\n"
   ]
  }
 ],
 "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
}