{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/8.clustering/Clustering_Sentence_Embeddings.ipynb)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zZWo_wE_U0i4"
   },
   "source": [
    "This notebook explores the use of SentenceBERT to generate representations of sequences (sentences, documents) and clustering those representations using K-means."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xsLXe22vZRyY"
   },
   "outputs": [],
   "source": [
    "!pip install sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SEvo_V2kZcf5"
   },
   "outputs": [],
   "source": [
    "# Get movies summaries and book titles to cluster\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/plot_summaries.txt\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/loc/dev.tsv -O book_titles.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cPnLGCaGSOGN"
   },
   "outputs": [],
   "source": [
    "from math import sqrt\n",
    "\n",
    "import numpy as np\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sklearn.cluster import KMeans\n",
    "from tqdm import tqdm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j04I_-NfZZEM"
   },
   "outputs": [],
   "source": [
    "def read_data(filename):\n",
    "    data = []\n",
    "    with open(filename, encoding=\"utf-8\") as file:\n",
    "        for line in file:\n",
    "            cols = line.rstrip().split(\"\\t\")\n",
    "            idd = cols[0]\n",
    "            summary = cols[1]\n",
    "            data.append((idd, summary))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_VviGH3nZiiH"
   },
   "outputs": [],
   "source": [
    "movies = read_data(\"plot_summaries.txt\")\n",
    "book_titles = read_data(\"book_titles.txt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the sentence embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SmChmfWtZmYc"
   },
   "outputs": [],
   "source": [
    "sentence_model = SentenceTransformer('sentence-transformers/all-distilroberta-v1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's try embedding a sentence. What is the shape of the embedding?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HacVXkwQaU9V"
   },
   "outputs": [],
   "source": [
    "embedding = sentence_model.encode(\"this is a sentence\")\n",
    "print(embedding.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YjWG2l73O5MB"
   },
   "outputs": [],
   "source": [
    "def cosine_similarity(one, two):\n",
    "  return np.dot(one,two) / (sqrt(np.dot(one,one)) * sqrt(np.dot(two,two)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embeddings(data, model):\n",
    "    X = []\n",
    "\n",
    "    # Get sentence embeddings for each doc\n",
    "    for idx, doc in tqdm(data):\n",
    "        embedding = model.encode(doc)\n",
    "        X.append(embedding)\n",
    "\n",
    "    return np.array(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_wi9-2aAay34"
   },
   "outputs": [],
   "source": [
    "def run_all(data, model, num_clusters=10):\n",
    "\n",
    "    embeddings = get_embeddings(data, model)\n",
    "\n",
    "    # Run K-means\n",
    "    kmeans = KMeans(n_clusters=num_clusters, random_state=0).fit(embeddings)\n",
    "\n",
    "    # For each cluster, print out the n documents closest to the cluster center\n",
    "    clusters = {}\n",
    "    for idx, label in enumerate(kmeans.labels_):\n",
    "        if label not in clusters:\n",
    "            clusters[label] = []\n",
    "        clusters[label].append((idx, data[idx][1]))\n",
    "\n",
    "    for label in clusters:\n",
    "        sims = {}\n",
    "        cluster_center = kmeans.cluster_centers_[label]\n",
    "        for idx, doc in clusters[label]:\n",
    "            sim = cosine_similarity(cluster_center, embeddings[idx])\n",
    "            sims[idx] = sim\n",
    "        for k, v in sorted(sims.items(), key=lambda item: item[1], reverse=True)[:5]:\n",
    "            print(k,\"%.3f\" % v, data[k][1])\n",
    "        \n",
    "        print()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "L_hOcnwzVgV9"
   },
   "source": [
    "# Book titles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OrddyPw1bHdG"
   },
   "outputs": [],
   "source": [
    "run_all(book_titles[:1000], sentence_model, num_clusters=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HxDLind6VSTy"
   },
   "source": [
    "# Movie summaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ngWb-mi8bLfy"
   },
   "outputs": [],
   "source": [
    "run_all(movies[:100], sentence_model, num_clusters=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-GgrP1OqVvU4"
   },
   "source": [
    "**Q1**: Play around with this method and vary the number of movies clustered, along with the number of clusters.  How would you rate the coherence and interepretability of these clusters? Try to label some of the clusters and discuss with your neighbors about the overall coherence."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}