LLM survey Responses Labelling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analysis of Survey Responses using LLMs and Clustering\n",
    "\n",
    "This notebook contains code that can label and cluster survey responses.\n",
    "\n",
    "It is designed to be ran in the Federated Data Platform.\n",
    "\n",
    "This notebook will explain teh process of annotating responses from one survey location. However, the code can easily be edited to handle more datasets."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Method\n",
    "\n",
    "### 1 - Setup and Labelling\n",
    "\n",
    "First, we will import the required modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from foundry.transforms import Dataset\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.cluster import AffinityPropagation\n",
    "import json\n",
    "import re\n",
    "import ast\n",
    "\n",
    "from language_model_service_api.languagemodelservice_api_embeddings_v3 import GenericEmbeddingsRequest\n",
    "from language_model_service_api.languagemodelservice_api_completion_v3 import GptChatCompletionRequest\n",
    "from language_model_service_api.languagemodelservice_api import ChatMessage, ChatMessageRole\n",
    "\n",
    "from palantir_models.models import GenericEmbeddingModel\n",
    "from palantir_models.models import OpenAiGptChatLanguageModel\n",
    "\n",
    "from string import Template\n",
    "\n",
    "class PromptTemplate(Template):\n",
    "    delimiter = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will import the required data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "responses = Dataset.get(\"<<Name of Dataset>>\").read_table()\n",
    "\n",
    "# Optional - load in multiple data sources.\n",
    "responses_2 = Dataset.get(\"<<Name of Dataset>>\").read_table()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define a prompt which can analyse a response and output the sentiment, alongside ideas for change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = PromptTemplate(\"\"\"\n",
    "You identify 'Sentiment' and 'Ideas for Change' from a survey response. \n",
    "\n",
    "You are to note on 'Ideas for Change' and 'Sentiment', alongside giving your 'Reasoning'.\n",
    "\n",
    "Each of these fields is a free text field. Sentiment must be one of 'Positive', 'Neutral', 'Poor' or 'Very Poor'.\n",
    "\n",
    "Your response must be valid json.\n",
    "\n",
    "Here are two examples of key topics extracted from survey responses.\n",
    "\n",
    "Example 1:\n",
    "{example_1_prompt}\n",
    "\n",
    "The output would be:\n",
    "{example_1_output}\n",
    "\n",
    "Example 2:\n",
    "{example_2_prompt}\n",
    "\n",
    "The output would be:\n",
    "{example_2_output}\n",
    "\n",
    "Here is the survey reponse you must label with topics:\n",
    "\n",
    "Survey Response:\n",
    "{survey_response}\n",
    "\n",
    "Within the 'Reasoning' section of your json output, think step by step to reason each topic. This are is for you to outline your thoughts, consider if you have missed any key topics, and check your responses.\n",
    "Then, score on 'Sentiment' and 'Ideas for Change'. Your response must be valid json.\n",
    "\n",
    "Your response must be in the format:\n",
    "\n",
    "{\"Reasoning\":\"xx\",\"Ideas for Change\": \"xx\", \"Sentiment\": \"xx\"}\n",
    "\n",
    "Do not put any personal identifiable information in your response. Include nothing else in your output.\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also can define a function which given a database row, will output a readable example that can be given to a prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def row_to_prompt(row,\n",
    "                  row_columns):\n",
    "    \n",
    "    string = ''\n",
    "    \n",
    "    for column, response in zip(row_columns, row[row_columns]):\n",
    "        new_string = f'{column.replace(\"_\",\" \")}\\n{response}\\n\\n'\n",
    "        string += new_string\n",
    "        \n",
    "    return string"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently, structured outputs are not supported within the FDP.\n",
    "\n",
    "We will use a cleaning prompt to clean json outputs if needed. That is defined here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_cleaning = \"\"\"Clean this json file: {json_file}. Only output a json file, include no other text. The json should be of the format: {\"Reasoning\":\"xx\",\"Short Term priorities\": \"xx\",\"Long Term Priorities\": \"xx\",\"Blockers\": \"xx\",\"Ideas for Change\": \"xx\"\"Sentiment\": \"xx\"}\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we select which columns we want to include in the examples given to our prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns_to_include_for_prompt = [\"<<Column 1 Name>>\", \"<<Column 2 Name>>\", ...]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We next define the examples we want to feed into out prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_1_prompt = row_to_prompt(responses.loc[0],columns_to_include_for_prompt)\n",
    "# You will need to give the example output\n",
    "example_1_output = {\"Reasoning\":\"xx\",\"Ideas for Change\": \"xx\", \"Sentiment\": \"xx\"}\n",
    "\n",
    "example_2_prompt = row_to_prompt(responses.loc[1],columns_to_include_for_prompt)\n",
    "# You will need to give the example output\n",
    "example_2_output = {\"Reasoning\":\"xx\",\"Ideas for Change\": \"xx\", \"Sentiment\": \"xx\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let us call an LLM and label each response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = OpenAiGptChatLanguageModel.get(\"GPT_4o\")\n",
    "\n",
    "LLM_responses = [example_1_output,example_2_output]\n",
    "length = len(reponses)\n",
    "\n",
    "print(\"Starting...\")\n",
    "for i in range(2, length):\n",
    "    \n",
    "    if i%100==0:\n",
    "        print(f\"Completed {i}/{length}\")\n",
    "    \n",
    "    formatted_prompt = prompt.safe_substitute(\n",
    "        example_1_prompt = example_1_prompt,\n",
    "        example_1_output = example_1_output,\n",
    "        example_2_prompt = example_2_prompt,\n",
    "        example_2_output = example_2_output,\n",
    "        survey_response = row_to_prompt(reponses.iloc[i],columns_to_include_for_prompt))\n",
    "    \n",
    "    response = model.create_chat_completion(GptChatCompletionRequest([ChatMessage(ChatMessageRole.USER, formatted_prompt)]))\n",
    "    raw_content = response.choices[0].message.content  # Extract the raw content\n",
    "    \n",
    "    try: \n",
    "        json_content = json.loads(raw_content)\n",
    "        LLM_responses.append(json_content)\n",
    "    \n",
    "    except:\n",
    "        \n",
    "        try:\n",
    "            raw_content = raw_content.replace(\"```\",\"\")\n",
    "            raw_content = raw_content.replace(\"json\",\"\")\n",
    "            raw_content = raw_content.replace(\"\\n\",\"\")\n",
    "            json_content = json.loads(raw_content)\n",
    "            LLM_responses.append(json_content)\n",
    "        except:\n",
    "            \n",
    "            try:\n",
    "                cleaning_prompt = cleaning_prompt.safe_substitute(json_file = raw_content)\n",
    "                response = model.create_chat_completion(GptChatCompletionRequest([ChatMessage(ChatMessageRole.USER, formatted_prompt)]))\n",
    "                raw_content = response.choices[0].message.content  # Extract the raw content\n",
    "                json_content = json.loads(raw_content)\n",
    "                LLM_responses.append(json_content)\n",
    "            except:\n",
    "                LLM_responses.append(f\"Error: {raw_content}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can save these responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = responses\n",
    "output[\"labels\"] = LLM_responses\n",
    "\n",
    "ideas_for_change = [x[\"Ideas for Change\"] for x in output[\"labels\"]]\n",
    "sentiments = [x[\"Sentiment\"] for x in output[\"labels\"]]\n",
    "\n",
    "output[\"Ideas_for_Change_LLM\"] = ideas_for_change\n",
    "output[\"Sentiment\"] = sentiments\n",
    "\n",
    "final_labels = Dataset.get(\"<<Dataset Name>>\")\n",
    "final_labels.write_table(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2 - LLM Topic Modelling\n",
    "\n",
    "We can continue to use a LLM for topic modelling. This has been shown to be effect, for example in this [paper](https://arxiv.org/pdf/2403.16248)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_ideas = ideas_for_change # you can add together ideas for change from other documents here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are going to use an LLM to label the 'ideas to change' with topics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "extraction_prompt = PromptTemplate(\"\"\"\n",
    "You are categorizing \"Ideas for Change\" from survey responses into topics.\n",
    "\n",
    "- The \"Current Topics\" are: {topics_}\n",
    "- The \"Ideas for Change\" are: {ideas_for_change}\n",
    "\n",
    "Instructions:\n",
    "1. Assign the \"Ideas for Change\" to one or more topics from the \"Current Topics\".\n",
    "2. If the appropriate topic does not exist in \"Current Topics\", create a new topic and include it in your response. Give your new topic a descriptive name.\n",
    "3. If the \"Ideas for Change\" applies to two topics, you may assign it to two topics.\n",
    "4. Do not assign more than two topics.\n",
    "\n",
    "Example:\n",
    "If the \"Current Topics\" are [\"Improved Communication\", \"Employee Wellbeing\"] and the \"Idea for Change\" is \"Install solar panels on the roof and improve communication between staff members.\" your response might be:\n",
    "[\"Sustainability\", \"Improved Communication\"]\n",
    "\n",
    "Your response must strictly be a Python list containing the relevant topics, such as:\n",
    "[\"Topic A\", \"Topic B\"]\n",
    "\n",
    "Do not include any additional text in your response.\n",
    "Do not label a topic as a \"New Topic\" - instead give it a descriptive name.\n",
    "\"\"\")\n",
    "\n",
    "cleaning_prompt = PromptTemplate(\"\"\"\n",
    "Clean this list of topics to be a flat python list. The list is: {response}.\n",
    "\n",
    "Your response must strictly be a flat python list containing the relevant topics, such as:\n",
    "[\"Topic A\", \"Topic B\"]\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We next call our LLM and apply topics to each idea."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "current_topics = [\"Improved Communication\"]\n",
    "\n",
    "all_responses = []\n",
    "failed_responses = []\n",
    "\n",
    "model = OpenAiGptChatLanguageModel.get(\"GPT_4o\")\n",
    "\n",
    "def add_topics(all_responses, topic_array, current_topics):\n",
    "    \n",
    "    all_responses.append(topic_array)\n",
    "    \n",
    "    for topic in topic_array:\n",
    "        if topic not in current_topics:\n",
    "            current_topics.append(topic)\n",
    "            \n",
    "    return all_responses, topic_array, current_topics\n",
    "\n",
    "def contains_nests(arr):\n",
    "    for element in arr:\n",
    "        if isinstance(element, list):\n",
    "            return True\n",
    "    return False\n",
    "    \n",
    "def flatten(array):\n",
    "        result = []\n",
    "        for item in array:\n",
    "            if isinstance(item, list):\n",
    "                result.extend(flatten(item))  # Recursively flatten nested lists\n",
    "            else:\n",
    "                result.append(item)\n",
    "        return result\n",
    "\n",
    "for i in range(len(all_ideas)):\n",
    "    \n",
    "    formatted_prompt = extraction_prompt.safe_substitute(\n",
    "            topics_ = current_topics,\n",
    "            ideas_for_change = all_ideas[i])\n",
    "    \n",
    "    response = model.create_chat_completion(GptChatCompletionRequest([ChatMessage(ChatMessageRole.USER, formatted_prompt)]))\n",
    "    raw_content = response.choices[0].message.content  # Extract the raw content\n",
    "    try:\n",
    "        topic_array = ast.literal_eval(raw_content)\n",
    "        if contains_nests(topic_array):\n",
    "            topic_array = flatten(topic_array)\n",
    "        all_responses, topic_array, current_topics = add_topics(all_responses, topic_array, current_topics)\n",
    "    except:\n",
    "        try:\n",
    "            formatted_cleaning_prompt = cleaning_prompt.safe_substitute(response=raw_content)\n",
    "            response = model.create_chat_completion(GptChatCompletionRequest([ChatMessage(ChatMessageRole.USER, formatted_prompt)]))\n",
    "            raw_content = response.choices[0].message.content  # Extract the raw content\n",
    "            topic_array = ast.literal_eval(raw_content)\n",
    "            \n",
    "            if contains_nests(topic_array):\n",
    "                topic_array = flatten(topic_array)\n",
    "            all_responses, topic_array, current_topics = add_topics(all_responses, topic_array, current_topics)\n",
    "            \n",
    "        except:\n",
    "            print(\"Failed at\",i)\n",
    "            print(\"with\", raw_content)\n",
    "            all_responses.append(\"ERROR: \")\n",
    "            failed_responses.append([i, raw_content]) \n",
    "    \n",
    "    if i%100==0:\n",
    "        print(f\"Completed {i}/{len(all_ideas)}\")\n",
    "        print(f\"Current Topics: {current_topics}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are various steps of cleaning required to reduce the number of errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cleaned_ideas = []\n",
    "cleaned_responses = []\n",
    "\n",
    "for idea in all_ideas:\n",
    "    if isinstance(idea, str):\n",
    "        cleaned_ideas.append(idea)\n",
    "    else:\n",
    "        cleaned_ideas.append(\"None\")\n",
    "    \n",
    "\n",
    "for response in all_responses:\n",
    "    if isinstance(response, list):\n",
    "        cleaned_responses.append(response)\n",
    "    else:\n",
    "        cleaned_responses.append([\"ERROR\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at how many topics there are, and save our first lot of data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"The number of topics are:\", len(current_topics))\n",
    "\n",
    "df = pd.DataFrame({\"Ideas for Change\": cleaned_ideas, \"Topics\": cleaned_responses})\n",
    "\n",
    "print(\"Modelling failed at:\", [f[0] for f in failed_responses])\n",
    "\n",
    "llm_topic_modeling = Dataset.get(\"llm_topic_modeling\")\n",
    "llm_topic_modeling.write_table(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now explore the topis in a bit more detail, and save another DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_df = {\"Topic\": [],\n",
    "           \"Count\": [],\n",
    "           \"Cluster\": []}\n",
    "\n",
    "for topic_id in current_topics:\n",
    "    \n",
    "    count = 0\n",
    "    ideas = []\n",
    "    \n",
    "    for topic_labels, idea in zip(all_responses,all_ideas):\n",
    "        if isinstance(topic_labels, list) and topic_id in topic_labels:\n",
    "            count += 1\n",
    "            ideas.append(idea)\n",
    "        \n",
    "    topic_df[\"Topic\"].append(topic_id)\n",
    "    topic_df[\"Count\"].append(count)\n",
    "    topic_df[\"Cluster\"].append(ideas)\n",
    "\n",
    "pd_topic_df = pd.DataFrame(topic_df).sort_values(by=\"Count\",ascending=False)\n",
    "\n",
    "llm_topic_counts = Dataset.get(\"llm_topic_counts\")\n",
    "llm_topic_counts.write_table(pd_topic_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3 - Clustering\n",
    "\n",
    "Using our labelled responses, we can use embedding and clustering to analyse the responses."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We expect that some responses will have multiple ideas for change. For that reason, we do analysis on the full responses to 'Ideas for Change', alongside responses split by a comma."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_ideas = ideas_for_change # you can add together ideas for change from other documents here.\n",
    "\n",
    "split_all_ideas = []\n",
    "cleaned_all_ideas = []\n",
    "\n",
    "for idea in all_ideas:\n",
    "    if not pd.isna(idea):\n",
    "        cleaned_all_ideas.append(idea)\n",
    "        split_all_ideas.extend(idea.split(\",\"))\n",
    "    else:\n",
    "        split_all_ideas.append(\"None\")\n",
    "        cleaned_all_ideas.append(\"None\")\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We alsop have to do some cleaning of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_ideas = [string for string in cleaned_all_ideas if string !=\"\"]  \n",
    "split_all_ideas = [string for string in split_all_ideas if string !=\"\"]  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We split the data into 'batches' - smaller chunks which can be passed into models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split_batch_locations = []\n",
    "batch_locations = []\n",
    "\n",
    "for i in range(len(split_all_ideas)):\n",
    "    if i % 100 == 0 and i != 0:\n",
    "        split_batch_locations.append(i)\n",
    "        \n",
    "for i in range(len(all_ideas)):\n",
    "    if i % 100 == 0 and i != 0:\n",
    "        batch_locations.append(i)\n",
    "        \n",
    "split_batches = np.split(split_all_ideas, split_batch_locations)\n",
    "batches = np.split(all_ideas, batch_locations)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we calculate 'embeddings'. This converts our strings of text into vectors. These vectors are essentially coordinates, and we can use these coordinates to see how similar different strings of text are and group them into clusters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = []\n",
    "split_embeddings = []\n",
    "\n",
    "model = GenericEmbeddingModel.get(\"text-embedding-ada-002\")\n",
    "print(\"Calculating Embeddings for split ideas...\")\n",
    "\n",
    "for n, batch in enumerate(split_batches):\n",
    "    try:\n",
    "        response = model.create_embeddings(GenericEmbeddingsRequest(inputs=list(batch)))\n",
    "        split_embeddings.extend((response.embeddings))\n",
    "    except:\n",
    "        print(\"Error at\",n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Calculating Embeddings for ideas...\")\n",
    "for n, batch in enumerate(batches):\n",
    "    try:\n",
    "        response = model.create_embeddings(GenericEmbeddingsRequest(inputs=list(batch)))\n",
    "        embeddings.extend((response.embeddings))\n",
    "    except:\n",
    "        print(\"Error at\",n)\n",
    "        \n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have embeddings, we want to do some 'clustering'. \n",
    "\n",
    "For this example we use a technique called [Affinity Propogation](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.AffinityPropagation.html#sklearn.cluster.AffinityPropagation).\n",
    "\n",
    "This has the advantage of working in non-flat geometry, and you do not need to state how many clusters you want it to generate.\n",
    "\n",
    "However, **we have not tested any other methods yet**.\n",
    "\n",
    "You can view other methods [here](https://scikit-learn.org/stable/modules/clustering.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = np.array(embeddings)\n",
    "clusters = AffinityPropagation(random_state=0).fit(embeddings)\n",
    "\n",
    "split_embeddings = np.array(split_embeddings)\n",
    "split_clusters = AffinityPropagation(random_state=0).fit(split_embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look at how many clusters there are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"When ideas are not split:\")\n",
    "print(f\"Number of Clusters: {split_clusters.labels_.max()}\")\n",
    "\n",
    "print(\"When ideas are not split:\")\n",
    "print(f\"Number of Clusters: {clusters.labels_.max()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using a more efficient function, we can automate the process of clustering, alongside getting an LLM to describe to us what appears in each cluster with a summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_summary_prompt = \"\"\"\n",
    "Given a list of topics, summarise the topics into a short, single sentence summary of approximately 10 words.\n",
    "\n",
    "You must only include information from the list of topics.\n",
    "Do not assume any additional information.\n",
    "\n",
    "The topics are: {topics}\n",
    "\"\"\"\n",
    "\n",
    "def cluster_analysis(embeddings, data, prompt):\n",
    "    \n",
    "    sentence_summaries = []\n",
    "    cluster_sizes = []\n",
    "    \n",
    "    model = OpenAiGptChatLanguageModel.get(\"GPT_4o\")\n",
    "    \n",
    "    embeddings = np.array(embeddings)\n",
    "    clusters = AffinityPropagation(random_state=1).fit(embeddings)\n",
    "    \n",
    "    print(f\"Number of Clusters: {clusters.labels_.max()}\")\n",
    "    \n",
    "    for cluster_id in range(clusters.labels_.max()):\n",
    "        \n",
    "        topic_cluster = []\n",
    "        \n",
    "        for label, string in zip(clusters.labels_, data):\n",
    "            if label == cluster_id:\n",
    "                topic_cluster.append(string)\n",
    "        \n",
    "        if cluster_id % 10 == 0:\n",
    "            print(f\"Analysing Cluster {cluster_id}/{clusters.labels_.max()}\")\n",
    "        response = model.create_chat_completion(\n",
    "            GptChatCompletionRequest([ChatMessage(ChatMessageRole.USER, prompt.format(topics=topic_cluster))]))\n",
    "\n",
    "        cluster_sizes.append(len(topic_cluster))\n",
    "        sentence_summaries.append(response.choices[0].message.content)\n",
    "    \n",
    "    print(\"Done\")\n",
    "    return cluster_sizes, sentence_summaries, clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can call the function like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_sizes, cluster_summaries, cluster_ids = cluster_analysis(embeddings,\n",
    "                                                all_ideas,\n",
    "                                                cluster_summary_prompt,\n",
    "                                                )\n",
    "\n",
    "split_cluster_sizes, split_cluster_summaries, split_cluster_ids = cluster_analysis(split_embeddings,\n",
    "                                                split_all_ideas,\n",
    "                                                cluster_summary_prompt,\n",
    "                                                )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's save our results to a DataFrame for further analysis later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "topics = []\n",
    "\n",
    "for cluster_id in range(cluster_ids.labels_.max()):\n",
    "    topic_cluster = []\n",
    "    for label, string in zip(cluster_ids.labels_, all_ideas):\n",
    "        if label == cluster_id:\n",
    "            topic_cluster.append(string)\n",
    "    topics.append(topic_cluster)\n",
    "\n",
    "print(len(topics))\n",
    "LLM_df = pd.DataFrame({\"Counts\": cluster_sizes, \"Summaries\": cluster_summaries, \"Cluster Content\":topics})\n",
    "\n",
    "# Sort the clusters by size, largest to smallest.\n",
    "LLM_df_sorted = LLM_df.sort_values(by = \"Counts\", ascending = False)\n",
    "\n",
    "clusters_summary = Dataset.get(\"clusters_summary\")\n",
    "clusters_summary.write_table(LLM_df_sorted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split_topics = []\n",
    "\n",
    "for cluster_id in range(split_cluster_ids.labels_.max()):\n",
    "    topic_cluster = []\n",
    "    for label, string in zip(split_cluster_ids.labels_, all_ideas):\n",
    "        if label == cluster_id:\n",
    "            topic_cluster.append(string)\n",
    "    split_topics.append(topic_cluster)\n",
    "\n",
    "LLM_df_split = pd.DataFrame({\"Counts\": split_cluster_sizes, \"Summaries\": split_cluster_summaries, \"Cluster Content\":split_topics})\n",
    "\n",
    "# Sort the clusters by size, largest to smallest.\n",
    "LLM_df_sorted_split = LLM_df_split.sort_values(by = \"Counts\", ascending = False)\n",
    "\n",
    "clusters_summary_split = Dataset.get(\"clusters_summary_split\")\n",
    "clusters_summary_split.write_table(LLM_df_sorted_split)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This concludes the example in this notebook. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}