Synapse/notebooks/text_clustering/standalone_text_clustering_OpenAI.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Copyright (c) Microsoft Corporation. All rights reserved.\n",
        "\n",
        "Licensed under the MIT License.\n",
        "\n",
        "# Clustering in Spark using OpenAI and KMeans\n",
        "\n",
        "## This cell configures the spark session - Do not change\n"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%configure -f\n",
        "{\n",
        "\"conf\": {\n",
        "     \"spark.rpc.message.maxSize\": 1024,\n",
        "     \"spark.kryoserializer.buffer.max\": \"1024m\"\n",
        "   }\n",
        "}"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-06-21T13:57:12.2233276Z",
              "execution_start_time": "2022-06-21T13:57:12.2229805Z",
              "livy_statement_state": "available",
              "queued_time": "2022-06-21T13:56:37.3772262Z",
              "session_id": 33,
              "session_start_time": "2022-06-21T13:56:37.4982917Z",
              "spark_pool": null,
              "state": "finished",
              "statement_id": -1
            },
            "text/plain": "StatementMeta(, 33, -1, Finished, Available)"
          },
          "metadata": {}
        }
      ],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## These are the parameters that need to be changed to your values"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# The input file name \n",
        "input_filename = 'abfss://[].dfs.core.windows.net/kaggle_videos/output/kaggle_videos_all_captions_single.csv/kaggle_videos_all_captions_single.csv'\n",
        "# The number of clusters - this can be automated or start with a guesstimate\n",
        "number_of_clusters = 50\n",
        "# The output directory where the output file will be written to\n",
        "output_directory = 'abfss://[].dfs.core.windows.net/kaggle_videos/output/'\n",
        "# The name of the output file\n",
        "output_filename = 'kaggle_videos_all_captions_text_clustered.csv'\n",
        "\n",
        "# The blob account url - https://[accountname].blob.core.windows.net\n",
        "account_url = \"https://[].blob.core.windows.net\"\n",
        "# The blob account name = [accountname]\n",
        "account_name = ''\n",
        "# The blob account key [iufquq34r423r2==] - used to generate a SAS key\n",
        "account_key = ''\n",
        "\n",
        "# The name of the primary ADLS share\n",
        "file_system_name=\"share\"\n",
        "# The directory folders where your files reside  \n",
        "directory_name='bbc'  # kaggle/videos/\n",
        "\n",
        "# If set to true generate a 3D scatterplot otherwise 2D\n",
        "SCATTER_PLOT_3D = True\n",
        "# If this is set to True then the Coalesce notebook will need to be run to merge the partition files into a single file\n",
        "LOW_MEMORY_MODE = True\n",
        "\n",
        "# Concept Graph - # Get top N most connected nodes   \n",
        "number_of_connected_nodes = 5\n",
        "\n",
        "# Azure SubscriptionId\n",
        "subscription_id=\"\"\n",
        "# AzureML Workspace Resource Group\n",
        "resource_group=\"\"\n",
        "# AzureML Workspace Name\n",
        "workspace_name=\"\"\n",
        "\n",
        "# OpenAI deployment values\n",
        "openai_api_key = \"\"\n",
        "openai_api_type = \"\"\n",
        "openai_api_base = \"\"\n",
        "openai_api_version = \"\"\n",
        "openai_deployment_id = \"\" "
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-06-21T13:57:16.2267479Z",
              "execution_start_time": "2022-06-21T13:57:16.0729048Z",
              "livy_statement_state": "available",
              "queued_time": "2022-06-21T13:56:37.4447783Z",
              "session_id": 33,
              "session_start_time": null,
              "spark_pool": "DataDiscovery",
              "state": "finished",
              "statement_id": 1
            },
            "text/plain": "StatementMeta(DataDiscovery, 33, 1, Finished, Available)"
          },
          "metadata": {}
        }
      ],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "tags": [
          "parameters"
        ]
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Track the Experiment in Azure ML"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "from azureml.core import Workspace, Experiment, Run\n",
        "import mlflow\n",
        "\n",
        "ws = Workspace(subscription_id=subscription_id, resource_group=resource_group, workspace_name=workspace_name)    \n",
        "mlflow.set_tracking_uri(ws.get_mlflow_tracking_uri())\n",
        "experiment_name = f\"({mssparkutils.runtime.context['notebookname']}_{str(mssparkutils.env.getJobId())})\"\n",
        "mlflow.set_experiment(experiment_name)\n",
        "mlflow.log_param(\"input_filename\", input_filename)\n",
        "mlflow.log_param(\"number_of_clusters\", number_of_clusters)\n",
        "mlflow.log_param(\"output_directory\", output_directory)\n",
        "mlflow.log_param(\"output_filename\", output_filename)\n",
        "mlflow.log_param(\"account_url\", account_url)\n",
        "mlflow.log_param(\"account_name\", account_name)\n",
        "mlflow.log_param(\"file_system_name\", file_system_name)\n",
        "mlflow.log_param(\"directory_name\", directory_name)\n",
        "mlflow.log_param(\"SCATTER_PLOT_3D\", SCATTER_PLOT_3D)\n",
        "mlflow.log_param(\"LOW_MEMORY_MODE\", LOW_MEMORY_MODE)\n",
        "mlflow.log_param(\"openai_deployment_id\", openai_deployment_id)\n",
        "mlflow.log_param(\"openai_api_version\", openai_api_version)\n",
        "params = {\n",
        "    \"sparkpool\": mssparkutils.runtime.context['sparkpool'],\n",
        "    \"workspace\": mssparkutils.runtime.context['workspace'],\n",
        "    \"notebookname\": mssparkutils.runtime.context['notebookname'],\n",
        "    \"isForPipeline\": mssparkutils.runtime.context['isForPipeline'],\n",
        "    \"pipelinejobid\": mssparkutils.runtime.context['pipelinejobid']\n",
        "}\n",
        "\n",
        "mlflow.log_params(params)\n",
        "mlflow.pyspark.ml.autolog()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "source": [
        "## Apply Data Preprocessing, OpenAI Embeddings, Kmeans clustering and PCA dimensionality reduction"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "from pyspark.ml.feature import HashingTF, IDF, Tokenizer, CountVectorizer, StopWordsRemover, PCA, RegexTokenizer\n",
        "from pyspark.ml.clustering import LDA, KMeans, BisectingKMeans\n",
        "from pyspark.ml import Pipeline\n",
        "from pyspark import SparkContext, SparkConf\n",
        "import sys\n",
        "from pyspark.sql.functions import udf\n",
        "from pyspark.sql.types import StringType\n",
        "from pyspark.sql import SparkSession\n",
        "import ntpath\n",
        "import os\n",
        "import numpy as np\n",
        "import openai\n",
        "\n",
        "from pyspark.ml.linalg import Vectors, VectorUDT\n",
        "from pyspark.sql.functions import udf,col\n",
        "\n",
        "from itertools import combinations\n",
        "from operator import itemgetter\n",
        "\n",
        "from graphframes import *\n",
        "from pyspark.sql.functions import monotonically_increasing_id, lit\n",
        "\n",
        "df = spark.read.load(input_filename, header=True, format='csv')\n",
        "\n",
        "def openai_embeddings(text):\n",
        "\n",
        "    from time import sleep\n",
        "    sleep(2) # Avoid rate limit for S0 tier\n",
        "    openai.api_key = openai_api_key\n",
        "    openai.api_type = openai_api_type\n",
        "    openai.api_base = openai_api_base\n",
        "    openai.api_version = openai_api_version\n",
        "    deployment_id = openai_deployment_id\n",
        "\n",
        "    embeddings = openai.Embedding.create(deployment_id=deployment_id,\n",
        "                                     input=text[:2000]) # Take first 2000 characters or calculate tokens\n",
        "    return Vectors.dense(embeddings['data'][0]['embedding'])\n",
        "\n",
        "# Vectorise processed text to OpenAI features\n",
        "udf_openai_vectoriser = udf(openai_embeddings, VectorUDT()) \n",
        "df = df.withColumn(\"OpenAI_features\", udf_openai_vectoriser(col(\"text\")))"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-06-21T14:01:05.2141837Z",
              "execution_start_time": "2022-06-21T13:57:16.3726956Z",
              "livy_statement_state": "available",
              "queued_time": "2022-06-21T13:56:37.4946798Z",
              "session_id": 33,
              "session_start_time": null,
              "spark_pool": "DataDiscovery",
              "state": "finished",
              "statement_id": 2
            },
            "text/plain": "StatementMeta(DataDiscovery, 33, 2, Finished, Available)"
          },
          "metadata": {}
        }
      ],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false
        }
      }
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "source": [
        "## Apply PCA and K-Means Clustering"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "# Apply PCA and Kmeans in a pipeline\n",
        "pca = PCA(k=20, inputCol=\"OpenAI_features\")\n",
        "pca.setOutputCol(\"features\")\n",
        "\n",
        "if SCATTER_PLOT_3D:\n",
        "  pca_2 = PCA(k=3, inputCol=\"features\")\n",
        "else:  \n",
        "  pca_2 = PCA(k=2, inputCol=\"features\")\n",
        "  \n",
        "pca_2.setOutputCol(\"pca_scatterplot_features\")\n",
        "\n",
        "kmeans = KMeans(k=number_of_clusters, seed=42, initMode=\"k-means||\", distanceMeasure=\"euclidean\")\n",
        "\n",
        "pipeline = Pipeline(stages=[ pca, kmeans, pca_2])\n",
        "model = pipeline.fit(df)\n",
        "df_coords = model.transform(df)"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Retrieve Data from Blob Storage"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "from azure.identity import DefaultAzureCredential\n",
        "from datetime import datetime, timedelta\n",
        "from azure.storage.blob import BlobServiceClient, generate_container_sas, BlobSasPermissions\n",
        "token_credential = DefaultAzureCredential()\n",
        "\n",
        "blob_service_client = BlobServiceClient(\n",
        "    account_url=account_url,\n",
        "    credential=token_credential\n",
        ")\n",
        "\n",
        "from azure.storage.filedatalake import DataLakeServiceClient, generate_directory_sas\n",
        "SAS = generate_directory_sas(\n",
        "        account_name=account_name,\n",
        "        file_system_name=file_system_name,\n",
        "        directory_name=directory_name,\n",
        "        credential=account_key,\n",
        "        permission=BlobSasPermissions(read=True),\n",
        "        expiry=datetime.utcnow() + timedelta(days=100))\n",
        "\n",
        "SAS_key  = \"?\" + SAS\n",
        "\n",
        "storage_path = os.path.join(account_url, file_system_name, directory_name)\n",
        "SAS_path = []\n",
        "\n",
        "def build_sas_path(row):\n",
        "    file_name = ntpath.basename(row)\n",
        "    return account_url + \"/\" + file_system_name + \"/\" + directory_name + \"/\" + file_name + SAS_key\n",
        "\n",
        "udf_build_sas_path = udf(build_sas_path, StringType())\n",
        "\n",
        "def get_X(row):\n",
        "    return str(row.values[0])\n",
        "\n",
        "def get_Y(row):\n",
        "    return str(row.values[1])\n",
        "\n",
        "def get_Z(row):\n",
        "    return str(row.values[2])\n",
        "\n",
        "def remove_separator(row):\n",
        "    # Clean up the separator from the string\n",
        "    return row.replace(\",\", \"\").replace(\"\\\\\", \"\")\n",
        "\n",
        "udf_get_X = udf(get_X, StringType())\n",
        "udf_get_Y = udf(get_Y, StringType())\n",
        "udf_get_Z = udf(get_Z, StringType())\n",
        "\n",
        "udf_remove_separator = udf(remove_separator, StringType())\n",
        "\n",
        "if SCATTER_PLOT_3D:\n",
        "    df_coords = df_coords.withColumn(\"blob_path\", udf_build_sas_path(df_coords.filename)).withColumn(\"X\", udf_get_X(df_coords.pca_scatterplot_features).cast('string')).withColumn(\"Y\", udf_get_Y(df_coords.pca_scatterplot_features).cast('string')).withColumn(\"Z\", udf_get_Z(df_coords.pca_scatterplot_features).cast('string')).withColumn(\"text\", udf_remove_separator(df.text))\n",
        "else:\n",
        "    df_coords = df_coords.withColumn(\"blob_path\", udf_build_sas_path(df_coords.filename)).withColumn(\"X\", udf_get_X(df_coords.pca_scatterplot_features).cast('string')).withColumn(\"Y\", udf_get_Y(df_coords.pca_scatterplot_features).cast('string')).withColumn(\"text\", udf_remove_separator(df.text))\n",
        "\n",
        "df_graph = df_coords\n",
        "df_coords = df_coords.drop('pca_features', 'pca_scatterplot_features', 'tokens', 'word_count_vector', 'pca_features', 'features', 'rawFeatures', 'stopWordsRemovedTokens', 'OpenAI_features')\n"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-06-21T14:01:07.1248785Z",
              "execution_start_time": "2022-06-21T14:01:05.3411899Z",
              "livy_statement_state": "available",
              "queued_time": "2022-06-21T13:56:37.5661166Z",
              "session_id": 33,
              "session_start_time": null,
              "spark_pool": "DataDiscovery",
              "state": "finished",
              "statement_id": 3
            },
            "text/plain": "StatementMeta(DataDiscovery, 33, 3, Finished, Available)"
          },
          "metadata": {}
        }
      ],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "collapsed": false
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Save the outcome with two operations (usual or coalesce)"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "if LOW_MEMORY_MODE:\n",
        "    df_coords.write.mode('overwrite').options(header='true').csv(os.path.join(output_directory, output_filename))\n",
        "else:\n",
        "    df1 = df_coords.filter((df_coords.blob_path != 'blob_path'))\n",
        "    df1.coalesce(1).write.mode('overwrite').options(header='true').csv(os.path.join(output_directory, output_filename))\n",
        "\n",
        "mlflow.pyspark.ml.mlflow.end_run()"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "execution_finish_time": "2022-06-21T14:01:12.5087704Z",
              "execution_start_time": "2022-06-21T14:01:07.2218654Z",
              "livy_statement_state": "available",
              "queued_time": "2022-06-21T13:56:37.6361295Z",
              "session_id": 33,
              "session_start_time": null,
              "spark_pool": "DataDiscovery",
              "state": "finished",
              "statement_id": 4
            },
            "text/plain": "StatementMeta(DataDiscovery, 33, 4, Finished, Available)"
          },
          "metadata": {}
        }
      ],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "outputs_hidden": false,
          "source_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Optional: Add Azure Cognitive Search"
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Add Search Parameters"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "# Azure Search Admin Key\n",
        "search_admin_key = \"\"\n",
        "# The name of the search service\n",
        "search_service_name = \"\"\n",
        "# The Azure Search Query Key\n",
        "search_query_key = \"\""
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "from synapse.ml.cognitive import *\n",
        "from pyspark.sql.functions import monotonically_increasing_id, lit\n",
        "\n",
        "df = df.drop(\"_c0\")\n",
        "\n",
        "(\n",
        "    df.withColumn(\"key\", monotonically_increasing_id().cast(\"string\"))\n",
        "    .withColumn(\"SearchAction\", lit(\"upload\"))\n",
        "    .writeToAzureSearch(\n",
        "        subscriptionKey=search_admin_key,\n",
        "        actionCol=\"SearchAction\",\n",
        "        serviceName=search_service_name,\n",
        "        indexName=experiment_name,  # Defaults to the notebook name\n",
        "        keyCol=\"key\",\n",
        "    )\n",
        ")"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Search the generated Azure Search Index"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "from IPython.display import display, Markdown\n",
        "\n",
        "term_to_search_for = \"covid\"\n",
        "\n",
        "url = \"https://{}.search.windows.net/indexes/{}/docs/search?api-version=2019-05-06\".format(\n",
        "    search_service_name, experiment_name\n",
        "\n",
        ")\n",
        "jdata = requests.post(url, json={\"search\": term_to_search_for}, headers={\"api-key\": search_query_key}).json()\n",
        "\n",
        "for doc in jdata['value']:\n",
        "    display(Markdown(f'**Search Score {doc[\"@search.score\"]}** Document {doc[\"filename\"]}'))\n",
        "    display(Markdown(f'{doc[\"text\"]}'))\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Implement Semantic Search"
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "1) [Enable Semantic Search](https://docs.microsoft.com/en-us/azure/search/semantic-search-overview#enable-semantic-search) on your search instance\n",
        "\n",
        "2) [Configure Semantic Search](https://docs.microsoft.com/en-us/azure/search/semantic-how-to-query-request?tabs=semanticConfiguration%2Cportal#create-a-semantic-configuration)"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "term_to_search_for = \"Whose thumb was fractured?\"\n",
        "\n",
        "url = \"https://{}.search.windows.net/indexes/{}/docs/search?api-version=2021-04-30-Preview\".format(\n",
        "    search_service_name, experiment_name\n",
        ")\n",
        "jdata = requests.post(url, json={\"search\": term_to_search_for, \"queryType\": \"semantic\", \"semanticConfiguration\": \"config\", \"queryLanguage\": \"en-us\", \"answers\": \"extractive|count-3\",\n",
        "\"captions\": \"extractive|highlight-true\",  \"highlightPreTag\": \"<mark>\",\"highlightPostTag\": \"</mark>\"}, headers={\"api-key\": search_query_key}).json()\n",
        "\n",
        "for doc in jdata['value']:\n",
        "    display(Markdown(f'**Search Score {doc[\"@search.score\"]}** **Search rerankerScore Score {doc[\"@search.rerankerScore\"]}** Document {doc[\"filename\"]}'))\n",
        "    display(Markdown(f'@search.captions {doc[\"@search.captions\"]}'))\n"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "source": [
        "# Optional: Build the Concept Graph"
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optional: Amend this section to build your concept graph"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "lst_text = df_graph.select('text').rdd.flatMap(lambda x: x).collect()\n",
        "\n",
        "lst_source_node = []\n",
        "lst_source_node_weight = []\n",
        "lst_source_node_label = []\n",
        "lst_target_node = []\n",
        "lst_target_node_weight = []\n",
        "lst_target_node_label = []\n",
        "lst_source_url = []\n",
        "lst_target_url = []\n",
        "lst_edge_weight = []\n",
        "lst_edge_colour_weight = []\n",
        "\n",
        "lst_g_nodes = []\n",
        "lst_g_edges = []\n",
        "\n",
        "dict_nodes = {}\n",
        "\n",
        "from itertools import combinations\n",
        "\n",
        "for i, row in enumerate(lst_text):\n",
        "    \n",
        "    combos = list(combinations(row, 2))\n",
        "   \n",
        "    for c in combos:\n",
        "        # First update edge weights\n",
        "        if (c[0] + \"_\" + c[1] not in dict_nodes) and (c[1] + \"_\" + c[0] not in dict_nodes):\n",
        "            dict_nodes[c[0] + \"_\" + c[1]] = 1 # initialise and create first combo\n",
        "        elif c[0] + \"_\" + c[1] in dict_nodes:\n",
        "            dict_nodes[c[0] + \"_\" + c[1]] += 1\n",
        "        elif c[1] + \"_\" + c[0] in dict_nodes:\n",
        "            dict_nodes[c[1] + \"_\" + c[0]] += 1\n",
        "\n",
        "    for c in combos:\n",
        "        lst_source_node.append(c[0])\n",
        "        lst_g_nodes.append((c[0], c[0]))\n",
        "        lst_target_node.append((c[1]))\n",
        "        lst_g_nodes.append((c[1], c[1]))\n",
        "        if c[0] + \"_\" + c[1] in dict_nodes:\n",
        "            lst_edge_weight.append(dict_nodes[c[0] + \"_\" + c[1]])\n",
        "            lst_source_node_weight.append(dict_nodes[c[0] + \"_\" + c[1]])\n",
        "            lst_target_node_weight.append(dict_nodes[c[0] + \"_\" + c[1]])\n",
        "            lst_g_edges.append((c[0],c[1], \"related\"))\n",
        "        else:\n",
        "            lst_edge_weight.append(dict_nodes[c[1] + \"_\" + c[0]])\n",
        "            lst_source_node_weight.append(dict_nodes[c[1] + \"_\" + c[0]])\n",
        "            lst_target_node_weight.append(dict_nodes[c[1] + \"_\" + c[0]])\n",
        "\n",
        "topn = dict(sorted(dict_nodes.items(), key = itemgetter(1), reverse = True)[:number_of_connected_nodes])\n",
        "\n",
        "# Assign edge weight colour\n",
        "for key in zip(lst_source_node, lst_target_node):\n",
        "    \n",
        "    if key[0] + \"_\" + key[1] in topn or key[1] + \"_\" + key[0] in topn:\n",
        "        lst_edge_colour_weight.append(\"red\")\n",
        "    else:\n",
        "        lst_edge_colour_weight.append(\"black\")\n",
        "\n",
        "\n",
        "# Create the Graph RDD\n",
        "columns = ['source', 'target', 'source_node_weight', 'target_node_weight', 'edge_weight', 'edge_colour']\n",
        "df_concept_graph = spark.createDataFrame(zip(lst_source_node, lst_target_node, lst_source_node_weight, lst_target_node_weight, lst_edge_weight, lst_edge_colour_weight), columns)\n",
        "  \n",
        "# Create a Vertex DataFrame with unique ID column \"id\"\n",
        "v = sqlContext.createDataFrame(lst_g_nodes, [\"id\", \"name\"])\n",
        "# Create an Edge DataFrame with \"src\" and \"dst\" columns\n",
        "e = sqlContext.createDataFrame(lst_g_edges, [\"src\", \"dst\", \"relationship\"])"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Show degree connectivity"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from graphframes import GraphFrame\n",
        "g = GraphFrame(v, e)\n",
        "\n",
        "# Query: Get in-degree of each vertex.\n",
        "df_degree = g.inDegrees\n",
        "df_degree.sort(['inDegree'], ascending=False).show()"
      ],
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.livy.statement-meta+json": {
              "spark_pool": null,
              "session_id": null,
              "statement_id": null,
              "state": "cancelled",
              "livy_statement_state": null,
              "queued_time": "2023-03-20T11:57:51.011691Z",
              "session_start_time": "2023-03-20T11:57:51.0662228Z",
              "execution_start_time": null,
              "execution_finish_time": "2023-03-20T11:57:51.8160924Z",
              "spark_jobs": null,
              "parent_msg_id": "aef23c4b-a479-4fee-9950-7c676f02dd9f"
            },
            "text/plain": "StatementMeta(, , , Cancelled, )"
          },
          "metadata": {}
        }
      ],
      "execution_count": 1,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run PageRank "
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Query: Count the number of \"follow\" connections in the graph.\n",
        "g.edges.filter(\"relationship = 'relationship'\").count()\n",
        "\n",
        "# Run PageRank algorithm, and show results.\n",
        "results = g.pageRank(resetProbability=0.01, maxIter=20)\n",
        "results.vertices.select(\"id\", \"pagerank\").show()"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "source": [
        "## Save the outcome with two operations (usual or coalesce)"
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "source": [
        "df_concept_graph.write.mode('overwrite').options(header='true').csv(os.path.join(output_directory, output_filename[:-4] + \"concept_graph.csv\"))"
      ],
      "outputs": [],
      "execution_count": null,
      "metadata": {}
    }
  ],
  "metadata": {
    "kernelspec": {
      "name": "synapse_pyspark",
      "language": "Python",
      "display_name": "Synapse PySpark"
    },
    "language_info": {
      "name": "python"
    },
    "kernel_info": {
      "name": "synapse_pyspark"
    },
    "description": null,
    "save_output": true,
    "synapse_widget": {
      "version": "0.1",
      "state": {}
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}