Notebooks/Spark-Example-25a-Mllib-Custom-Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BLFOWOTbMKGa"
   },
   "source": [
    "## Creating Custom Transformers For Pyspark MlLib\n",
    "Contributor: Tommy Lin (October 20th, 2024)\n",
    "\n",
    "### Overview\n",
    "Transformers are tools used to convert data into a desired format, usually for machine learning algorithms. Data from an input column is transformed and placed into an output column. Essentially, they leverage the Pyspark DataFrame's built-in function withColumn() to transform a column(s) in the dataset.\n",
    "\n",
    "### Common Built-in MLLib Transformers\n",
    "- Tokenizer\n",
    "- StopWordRemover\n",
    "- StandardScalar\n",
    "- CountVectorizer\n",
    "\n",
    "### Key Parts of a Transformer\n",
    "All these transformers are extended from the PySpark MLLib Transformer class. They also all have a <strong>transform()</strong> function, hence the name Transformer. When creating our own Transformer, we need to create a <strong>transform()</strong> function. Some transformers also have a fit() function along with the transform() function, though it is not required."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IGugg0jhMKGe"
   },
   "source": [
    "### Built-in MLLib Transformers vs Custom Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "GMULCJtvMKGe"
   },
   "outputs": [],
   "source": [
    "from pyspark import SparkContext\n",
    "from pyspark.sql import SparkSession\n",
    "\n",
    "sc = SparkContext(\"local\")\n",
    "spark = SparkSession.builder.getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HbEt4y1KMKGg"
   },
   "source": [
    "#### Tokenizer\n",
    "The built-in MLLib Tokenizer splits strings by \" \" and also converts the string to lovercase. For our custom Tokenizer, we will convert the strings in the input column to lowercase, then split the strings by \" \" and puts the results in the output column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jlJgYbFAMKGg",
    "outputId": "96f3c601-79db-4ff7-efef-426b262cdb83"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------------+------------------------------+\n",
      "|text                    |tokenized_text                |\n",
      "+------------------------+------------------------------+\n",
      "|He went to the store.   |[he, went, to, the, store.]   |\n",
      "|Dogs like to fight cats.|[dogs, like, to, fight, cats.]|\n",
      "|When did you do that?   |[when, did, you, do, that?]   |\n",
      "+------------------------+------------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.ml.feature import Tokenizer\n",
    "\n",
    "# Create array of text\n",
    "text = [('He went to the store.',), ('Dogs like to fight cats.',), ('When did you do that?',)]\n",
    "df = spark.createDataFrame(text, ['text'])\n",
    "\n",
    "# Tokenize the sentences into list of words using the built-in Tokenizer\n",
    "# inputCol is the column we want to convert\n",
    "# outputCol is where we will store the converted inputs\n",
    "tokenizer = Tokenizer(inputCol=\"text\", outputCol=\"tokenized_text\")\n",
    "tokenizer.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "p7NqwfqRMKGg",
    "outputId": "88319c25-f56c-4d01-fa8d-99c2da817c05"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------------+------------------------------+\n",
      "|text                    |tokenized_text                |\n",
      "+------------------------+------------------------------+\n",
      "|He went to the store.   |[he, went, to, the, store.]   |\n",
      "|Dogs like to fight cats.|[dogs, like, to, fight, cats.]|\n",
      "|When did you do that?   |[when, did, you, do, that?]   |\n",
      "+------------------------+------------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.ml import Transformer\n",
    "import pyspark.sql.functions as F\n",
    "\n",
    "class CustomTokenizer(Transformer):\n",
    "    def __init__(self, inputCol, outputCol):\n",
    "        super().__init__()\n",
    "        self.inputCol = inputCol\n",
    "        self.outputCol = outputCol\n",
    "\n",
    "    # Transform method that converts each sentence to tokens\n",
    "    # 1. Takes the specified input column\n",
    "    # 2. performs the operations on the column\n",
    "    # 3. puts the new values into an output column\n",
    "    def _transform(self, df):\n",
    "        return df.withColumn(self.outputCol, F.split(F.lower(F.col(self.inputCol)), \" \"))\n",
    "\n",
    "\n",
    "# Tokenizes data using custom Tokenizer\n",
    "# Same parameters inputCol and outputCol, just like the built-in tokenizer\n",
    "custom_tokenizer = CustomTokenizer(inputCol=\"text\", outputCol=\"tokenized_text\")\n",
    "custom_tokenizer.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WA8mAVSWMKGh"
   },
   "source": [
    "#### VectorAssembler\n",
    "The built-in VectorAssembler transformer groups several column values into one vector. For example, with columns \"age\", \"height\", and \"weight\", the VectorAssembler will group these features into an array [\"age\", \"height\", \"weight\"]. We can use the F.array function in our own transform() method to emulate this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "USXsIdkWMKGh",
    "outputId": "41d6ceb7-9130-4905-c31e-a83bc66aac09"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---+---------+\n",
      "|x1 |x2 |features |\n",
      "+---+---+---------+\n",
      "|1.0|2.4|[1.0,2.4]|\n",
      "|2.9|5.0|[2.9,5.0]|\n",
      "|3.3|1.0|[3.3,1.0]|\n",
      "|4.8|2.0|[4.8,2.0]|\n",
      "|5.0|4.6|[5.0,4.6]|\n",
      "|6.0|1.2|[6.0,1.2]|\n",
      "+---+---+---------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "\n",
    "x = [(1.0, 2.4),\n",
    "     (2.9, 5.0),\n",
    "     (3.3, 1.0),\n",
    "     (4.8, 2.0),\n",
    "     (5.0, 4.6),\n",
    "     (6.0, 1.2)]\n",
    "\n",
    "df = spark.createDataFrame(x, [\"x1\", \"x2\"])\n",
    "\n",
    "# Use VectorAssembler to combine columns into a feature array\n",
    "assembler = VectorAssembler(inputCols=[\"x1\", \"x2\"], outputCol=\"features\")\n",
    "assembler.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aS590zXoMKGh",
    "outputId": "7a5a1421-9db3-4354-8b4b-2e5d7e1c9185"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---+----------+\n",
      "|x1 |x2 |features  |\n",
      "+---+---+----------+\n",
      "|1.0|2.4|[1.0, 2.4]|\n",
      "|2.9|5.0|[2.9, 5.0]|\n",
      "|3.3|1.0|[3.3, 1.0]|\n",
      "|4.8|2.0|[4.8, 2.0]|\n",
      "|5.0|4.6|[5.0, 4.6]|\n",
      "|6.0|1.2|[6.0, 1.2]|\n",
      "+---+---+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.ml import Transformer\n",
    "import pyspark.sql.functions as F\n",
    "from pyspark.ml.linalg import VectorUDT\n",
    "\n",
    "class CustomVectorAssembler(Transformer):\n",
    "    def __init__(self, inputCols, outputCol):\n",
    "        super().__init__()\n",
    "        self.inputCols = inputCols\n",
    "        self.outputCol = outputCol\n",
    "\n",
    "    # Transform method that collects each column value into a feature array\n",
    "    # Takes all the input columns we want to group together\n",
    "    # puts the column values into arrays using F.array()\n",
    "    # Stores the results in the output column\n",
    "    def _transform(self, df):\n",
    "        return df.withColumn(self.outputCol, F.array(*self.inputCols))\n",
    "\n",
    "\n",
    "# Use custom Vector Assembler to assemble features into an array\n",
    "assembler = CustomVectorAssembler(inputCols=[\"x1\", \"x2\"], outputCol=\"features\")\n",
    "assembler.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xB1yzPU-MKGi"
   },
   "source": [
    "### Custom Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D5-Eum5lMKGi"
   },
   "source": [
    "#### StandardScalar + Addition\n",
    "We can also add our own functionality to our Transformers. What if we wanted to add a number to our data after scaling it? To do this, we can create a fit() method to calculate the mean and standard deviation, then use transform() to scale the values. After, we can add our desired value and output the results to the output column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Mbs0D1n7MKGi",
    "outputId": "d79f5340-65a7-4c2e-88de-3524a61c99eb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----+-------------------+------------------+\n",
      "|value|scaled             |scaled+addition   |\n",
      "+-----+-------------------+------------------+\n",
      "|1.0  |-1.4605934866804429|8.539406513319557 |\n",
      "|2.0  |-1.0954451150103321|8.904554884989668 |\n",
      "|3.0  |-0.7302967433402214|9.269703256659778 |\n",
      "|4.0  |-0.3651483716701107|9.634851628329889 |\n",
      "|5.0  |0.0                |10.0              |\n",
      "|6.0  |0.3651483716701107 |10.365148371670111|\n",
      "|7.0  |0.7302967433402214 |10.730296743340222|\n",
      "|8.0  |1.0954451150103321 |11.095445115010332|\n",
      "|9.0  |1.4605934866804429 |11.460593486680443|\n",
      "+-----+-------------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyspark.ml import Transformer\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "# Custom Transformer that scales a column on numbers\n",
    "class CustomStandardScaler(Transformer):\n",
    "\n",
    "    # inputCol and outputCol are required parameters\n",
    "    # addition is optional, default value is 0, which makes this a regular StandardScaler\n",
    "    def __init__(self, inputCol, outputCol, addition=0):\n",
    "        super().__init__()\n",
    "        self.inputCol = inputCol\n",
    "        self.outputCol = outputCol\n",
    "        self.addition = addition\n",
    "\n",
    "    def fit(self, df):\n",
    "        # Calculate mean and stddev for the input column\n",
    "        stats = df.select(F.mean(F.col(self.inputCol)).alias(\"mean\"),\n",
    "                          F.stddev(F.col(self.inputCol)).alias(\"stddev\"))\\\n",
    "                    .first()\n",
    "\n",
    "        # Set mean and std of column to use later\n",
    "        # we must run fit() before running transform()\n",
    "        self.mean = stats[\"mean\"]\n",
    "        self.stddev = stats[\"stddev\"]\n",
    "\n",
    "    def _transform(self, df):\n",
    "        # Scale the column using our calculated mean and sd and add the addition value\n",
    "        # Use the values from the inputCol, mean, and std to scale\n",
    "        # Add desired addition value\n",
    "        # Store the results in the output column\n",
    "        return df.withColumn(self.outputCol, (F.col(self.inputCol) - self.mean) / self.stddev + self.addition)\n",
    "\n",
    "\n",
    "# Create dataframe of numbers\n",
    "x = [(1.0,), (2.0,), (3.0,), (4.0,), (5.0,), (6.0,), (7.0,), (8.0,), (9.0,)]\n",
    "df = spark.createDataFrame(x, [\"value\"])\n",
    "\n",
    "# Use custom Standard Scaler with no addition\n",
    "scaler = CustomStandardScaler(inputCol=\"value\", outputCol=\"scaled\")\n",
    "scaler.fit(df)\n",
    "df = scaler.transform(df)\n",
    "\n",
    "# Use custom Standard Scaler with addition\n",
    "scaler2 = CustomStandardScaler(inputCol=\"value\", outputCol=\"scaled+addition\", addition=10)\n",
    "scaler2.fit(df)\n",
    "scaler2.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_hBmG-WUMKGi"
   },
   "source": [
    "#### Lemmatizer\n",
    "What if we wanted to lemmatize our text?\n",
    "MLLib doesn't have a built-in lemmatizer transformer we can use, but we can create our own.\n",
    "By creating our own, we can choose when to lemmatize the text. We can lemmatize either the string of text, or an array of tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ph5lt0kiMKGi",
    "outputId": "c2a3d802-e971-4d9b-ebc3-a01bcdf83f59"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------+------------------+\n",
      "|text                 |lemmatized        |\n",
      "+---------------------+------------------+\n",
      "|Bob is running       |Bob be run        |\n",
      "|Dreaming about eating|dream about eat   |\n",
      "|He went to the store |he go to the store|\n",
      "+---------------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import spacy\n",
    "from pyspark.ml import Transformer, Pipeline\n",
    "from pyspark.sql import functions as F\n",
    "from pyspark.sql.types import StringType, ArrayType\n",
    "\n",
    "# Transformer for Lemmatizing Strings\n",
    "# Example: \"I am running\" => \"I be run\"\n",
    "class CustomLemmatizerString(Transformer):\n",
    "    def __init__(self, inputCol, outputCol):\n",
    "        super().__init__()\n",
    "        self.inputCol = inputCol\n",
    "        self.outputCol = outputCol\n",
    "\n",
    "        # Initializes language from spacy\n",
    "        self.lang = spacy.load(\"en_core_web_sm\")\n",
    "\n",
    "    # Helper function that lemmatizes the words in a given text string\n",
    "    def lemmatize(self, text):\n",
    "        # Uses the spacy library to convert our string into documents\n",
    "        document = self.lang(text)\n",
    "\n",
    "        # Lemmatize each word in the document and returns a string representation of the lemmatized tokens\n",
    "        return ' '.join(token.lemma_ for token in document)\n",
    "\n",
    "    # Transform function, uses helper function\n",
    "    def _transform(self, df):\n",
    "        # Create a User-Defined function that lemmatizes a dataframe column\n",
    "        # uses the lemmatize function we created above and specifies the result type\n",
    "        lemmatize_udf = F.udf(self.lemmatize, StringType())\n",
    "\n",
    "        # Use our UDF to Lemmatize the input column\n",
    "        # Takes\n",
    "        return df.withColumn(self.outputCol, lemmatize_udf(df[self.inputCol]))\n",
    "\n",
    "\n",
    "# Create array of text\n",
    "text = [('Bob is running',), ('Dreaming about eating',), ('He went to the store',)]\n",
    "df = spark.createDataFrame(text, ['text'])\n",
    "\n",
    "# Use custom lemmatizer transformer on string\n",
    "lemmatizer = CustomLemmatizerString(inputCol=\"text\", outputCol=\"lemmatized\")\n",
    "lemmatizer.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_Ycw55ivMKGj",
    "outputId": "d7f525b4-bb84-47da-9794-4878c760c2b5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------+--------------------------+------------------------+\n",
      "|text                 |tokens                    |lemmatized              |\n",
      "+---------------------+--------------------------+------------------------+\n",
      "|Bob is running       |[bob, is, running]        |[bob, be, run]          |\n",
      "|Dreaming about eating|[dreaming, about, eating] |[dream, about, eat]     |\n",
      "|He went to the store |[he, went, to, the, store]|[he, go, to, the, store]|\n",
      "+---------------------+--------------------------+------------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Transformer for Lemmatizing an array of tokens\n",
    "# Example: [\"I\", \"am\", \"running\"] => [\"I\", \"be\", \"run\"]\n",
    "class CustomLemmatizerTokens(Transformer):\n",
    "    def __init__(self, inputCol, outputCol):\n",
    "        super().__init__()\n",
    "        self.inputCol = inputCol\n",
    "        self.outputCol = outputCol\n",
    "\n",
    "        # Initializes language from spacy\n",
    "        self.lang = spacy.load(\"en_core_web_sm\")\n",
    "\n",
    "    # Helper function that lemmatizes the words in a given text string\n",
    "    def lemmatize(self, tokens):\n",
    "        # Converts tokens into a string and creates spacy doc objject\n",
    "        text = self.lang(\" \".join(tokens))\n",
    "\n",
    "        # Lemmatizes words and returns as array\n",
    "        return [word.lemma_ for word in text]\n",
    "\n",
    "    # Transform function, uses helper function\n",
    "    def _transform(self, df):\n",
    "        # Create a User-Defined function that lemmatizes a dataframe column\n",
    "        # Uses the lemmatize method we created above, and specifies array result type\n",
    "        lemmatize_udf = F.udf(self.lemmatize, ArrayType(StringType()))\n",
    "\n",
    "        # Use our UDF to Lemmatize the input column\n",
    "        # Takes in a column containing arrays of words\n",
    "        # Lemmatizes each word of arrays in the input column\n",
    "        # Stores the arrays of lemmatized words in the output column\n",
    "        return df.withColumn(self.outputCol, lemmatize_udf(df[self.inputCol]))\n",
    "\n",
    "\n",
    "# Create array of text\n",
    "text = [('Bob is running',), ('Dreaming about eating',), ('He went to the store',)]\n",
    "df = spark.createDataFrame(text, ['text'])\n",
    "\n",
    "# Create my custom transformers\n",
    "tokenizer = CustomTokenizer(inputCol=\"text\", outputCol=\"tokens\")\n",
    "lemmatizer = CustomLemmatizerTokens(inputCol=\"tokens\", outputCol=\"lemmatized\")\n",
    "\n",
    "# Our custom Transformers can also be used in a Pipeline like built-in Transformers\n",
    "# Our lemmatizer Transformer can be used in our machine learning preprocessing pipeline\n",
    "\n",
    "# Create a pipeline and fit/transform data\n",
    "pipeline = Pipeline(stages=[tokenizer, lemmatizer])\n",
    "model = pipeline.fit(df)\n",
    "model.transform(df).show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gcexmUhnMKGj"
   },
   "source": [
    "## Notes for Custom Transformers\n",
    "- Must implement a _transform function in your custom transformer class\n",
    "- Must inherit from base MLLib Transformer class\n",
    "- Leverage DataFrame operations\n",
    "- Can be used in pipelines just as a built-in transformer\n",
    "- You can be as creative as possible!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "spark",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}