notebook/query_recommendation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e34cfb2b-b46d-4546-be3b-38b10b70a17e",
   "metadata": {},
   "source": [
    "# Query Recommendatation (Tag-to-Tag Similarity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "16b14b03-7010-4430-b8c9-906c9bdab2ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3eaf7699-30ee-4bc9-9b95-40589efc46f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mwe_dataset = load_dataset(\"seungheondoh/musical-word-embedding\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "05c6b255-ddf0-490a-be32-a139605c5ae8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_list = [i[\"content\"] for i in mwe_dataset['tag']]\n",
    "tag_vector = torch.tensor([i[\"vector\"] for i in mwe_dataset['tag']])\n",
    "tag2vector = {token:embs for token, embs in zip(tag_list, tag_vector)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "641f52e3-6589-4fb9-aeb2-c3f149db712e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tag_vector = torch.nn.functional.normalize(tag_vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "12974f6b-17d0-4bd7-b3e8-5c4fd78e196c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_recommnedation(user_query, topk=10):\n",
    "    token_list = [i.strip() for i in user_query.split()]\n",
    "    query_vector = torch.stack([tag2vector[i] for i in token_list])\n",
    "    query_vector = torch.nn.functional.normalize(query_vector)\n",
    "    if query_vector.size(0) > 1:\n",
    "        query_vector = query_vector.mean(0, True)\n",
    "    score_matrix = query_vector @ tag_vector.T\n",
    "    sorted_indices = torch.flip(torch.argsort(score_matrix, dim=1), dims=[1])\n",
    "    \n",
    "    score_matrix = score_matrix.squeeze(0)\n",
    "    sorted_indices = sorted_indices.squeeze(0)\n",
    "    \n",
    "    top5_indices = sorted_indices[:10]\n",
    "    results = []\n",
    "    for i in top5_indices:\n",
    "        results.append({\n",
    "            \"entity\": tag_list[i],\n",
    "            \"score\": float(score_matrix[i]),\n",
    "        })\n",
    "    return pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3865f375-8853-4311-86f7-45c8cea2382f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>entity</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>summer</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>summery</td>\n",
       "      <td>0.859929</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>happy</td>\n",
       "      <td>0.814526</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>carefree</td>\n",
       "      <td>0.778677</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>fun</td>\n",
       "      <td>0.770603</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>hangingout</td>\n",
       "      <td>0.767396</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>celebratory</td>\n",
       "      <td>0.766314</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>bright</td>\n",
       "      <td>0.766003</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>joyous</td>\n",
       "      <td>0.761268</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>cheerful</td>\n",
       "      <td>0.758519</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        entity     score\n",
       "0       summer  1.000000\n",
       "1      summery  0.859929\n",
       "2        happy  0.814526\n",
       "3     carefree  0.778677\n",
       "4          fun  0.770603\n",
       "5   hangingout  0.767396\n",
       "6  celebratory  0.766314\n",
       "7       bright  0.766003\n",
       "8       joyous  0.761268\n",
       "9     cheerful  0.758519"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_recommnedation(\"summer\", topk=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0748207f-551e-4dfa-a3cf-56005a4dccab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>entity</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>house</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>club</td>\n",
       "      <td>0.825078</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>deephouse</td>\n",
       "      <td>0.808028</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>electro</td>\n",
       "      <td>0.788683</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>techno</td>\n",
       "      <td>0.787671</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>progressivehouse</td>\n",
       "      <td>0.763063</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>electronica</td>\n",
       "      <td>0.752875</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>electrohouse</td>\n",
       "      <td>0.751341</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>trance</td>\n",
       "      <td>0.748585</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>dance</td>\n",
       "      <td>0.739289</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             entity     score\n",
       "0             house  1.000000\n",
       "1              club  0.825078\n",
       "2         deephouse  0.808028\n",
       "3           electro  0.788683\n",
       "4            techno  0.787671\n",
       "5  progressivehouse  0.763063\n",
       "6       electronica  0.752875\n",
       "7      electrohouse  0.751341\n",
       "8            trance  0.748585\n",
       "9             dance  0.739289"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_recommnedation(\"house\", topk=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9d732cd8-395c-46a0-a2b5-3422c8a1b6a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>entity</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>jazz</td>\n",
       "      <td>1.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>contemporaryjazz</td>\n",
       "      <td>0.788583</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>jazzinstrument</td>\n",
       "      <td>0.784430</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>vocaljazz</td>\n",
       "      <td>0.774813</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>smoothjazz</td>\n",
       "      <td>0.758968</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>postbop</td>\n",
       "      <td>0.758877</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>fusion</td>\n",
       "      <td>0.739702</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>jazzpiano</td>\n",
       "      <td>0.727406</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>jazzy</td>\n",
       "      <td>0.721525</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>swing</td>\n",
       "      <td>0.721187</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             entity     score\n",
       "0              jazz  1.000000\n",
       "1  contemporaryjazz  0.788583\n",
       "2    jazzinstrument  0.784430\n",
       "3         vocaljazz  0.774813\n",
       "4        smoothjazz  0.758968\n",
       "5           postbop  0.758877\n",
       "6            fusion  0.739702\n",
       "7         jazzpiano  0.727406\n",
       "8             jazzy  0.721525\n",
       "9             swing  0.721187"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_recommnedation(\"jazz\", topk=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa9413e-e8c9-4b22-a396-58b0ec84e608",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}