updated asessment and solutions notebooks

notebooks/assesment.ipynb

@@ -292,39 +292,32 @@
    "source": [
     "# 9. Bonus 2: Transformer Embeddings UDF\n",
     "# ======================================\n",
-    "# We'll demonstrate a simple approach using a lightweight transformer model to embed passenger names.\n",
-    "# This is optional, but shows advanced usage of Spark UDFs.\n",
-    "\n",
-    "# Requirements: e.g. \"transformers\" or \"sentence-transformers\" in your environment.\n",
-    "# from transformers import pipeline\n",
-    "# embedding_pipeline = pipeline(\"feature-extraction\", model=\"distilbert-base-uncased\")\n",
-    "# OR\n",
-    "# from sentence_transformers import SentenceTransformer\n",
-    "# model = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
-    "\n",
-    "# 9.1 TODO: import / load the model/pipeline\n",
-    "# e.g.\n",
-    "# from transformers import pipeline\n",
-    "# embedding_pipeline = pipeline(\"feature-extraction\", model=\"distilbert-base-uncased\")\n",
-    "\n",
-    "# 9.2 Define a Python function that takes a passenger name (string) -> returns a list of floats\n",
-    "\n",
-    "# def get_name_embedding(name: str) -> List[float]:\n",
-    "#     # TODO: use embedding_pipeline or model to produce an embedding\n",
-    "#     # embedding = ?\n",
-    "#     # NOTE: verify shape (embedding might be list of lists)\n",
-    "#     return ???\n",
-    "\n",
-    "# 9.3 Wrap that function in a PySpark UDF\n",
-    "# from pyspark.sql.functions import udf\n",
-    "# from pyspark.sql.types import ArrayType, FloatType\n",
-    "# udf_get_name_embedding = udf(get_name_embedding, ArrayType(FloatType()))\n",
-    "\n",
-    "# 9.4 Apply the UDF to create a new column 'NameEmbedding' in spark_merged_clean\n",
-    "# spark_embedded = spark_merged_clean.withColumn(\"NameEmbedding\", udf_get_name_embedding(F.col(\"Name\")))\n",
-    "\n",
-    "# spark_embedded.select(\"Name\", \"NameEmbedding\").show(truncate=False)\n"
+    "\n",
+    "from sentence_transformers import SentenceTransformer\n",
+    "from pyspark.sql.functions import udf\n",
+    "from pyspark.sql.types import ArrayType, FloatType\n",
+    "\n",
+    "# Load the pre-trained MiniLM sentence transformer model\n",
+    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
+    "\n",
+    "# Define a UDF to compute the embeddings\n",
+    "def compute_embedding(text):\n",
+    "    '''\n",
+    "    Your function goes here\n",
+    "    '''\n",
+    "    pass\n",
+    "\n",
+    "# Register the UDF in Spark\n",
+    "embedding_udf = None #Replace with your udf\n",
+    "\n",
+    "# Apply the UDF to compute embeddings for each document\n",
+    "df_with_embeddings = spark_merged_clean.withColumn('mini-lm-vectors', '...')"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

notebooks/solutions.ipynb

@@ -265,36 +265,26 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Example imports (make sure 'transformers' is installed)\n",
-    "from transformers import pipeline\n",
-    "embedding_pipeline = pipeline(\"feature-extraction\", model=\"distilbert-base-uncased\")\n",
-    "\n",
-    "# Example function to get the name embedding\n",
-    "def get_name_embedding(name: str):\n",
-    "    # The pipeline will return a list of lists of floats.\n",
-    "    # Typically shape: (1, sequence_length, hidden_size).\n",
-    "    # We'll take the first token or perhaps average them.\n",
-    "    output = embedding_pipeline(name)\n",
-    "    # output[0] is shape [sequence_length, hidden_size]\n",
-    "    # let's do a simple average across the sequence dimension:\n",
-    "    token_embeddings = output[0]\n",
-    "    # average across tokens:\n",
-    "    mean_embedding = [float(sum(x) / len(x)) for x in zip(*token_embeddings)]\n",
-    "    return mean_embedding\n",
-    "\n",
-    "# Convert this Python function to a Spark UDF\n",
+    "from sentence_transformers import SentenceTransformer\n",
     "from pyspark.sql.functions import udf\n",
     "from pyspark.sql.types import ArrayType, FloatType\n",
     "\n",
-    "udf_get_name_embedding = udf(get_name_embedding, ArrayType(FloatType()))\n",
+    "# Load the pre-trained MiniLM sentence transformer model\n",
+    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
     "\n",
-    "# Apply it to add a new column\n",
-    "spark_embedded = spark_merged_clean.withColumn(\n",
-    "    \"NameEmbedding\",\n",
-    "    udf_get_name_embedding(F.col(\"Name\"))\n",
-    ")\n",
+    "# Define a UDF to compute the embeddings\n",
+    "def compute_embedding(text):\n",
+    "    return model.encode(text).tolist()\n",
+    "\n",
+    "# Register the UDF in Spark\n",
+    "embedding_udf = udf(compute_embedding, ArrayType(FloatType()))\n",
+    "\n",
+    "# Apply the UDF to compute embeddings for each document\n",
+    "df_with_embeddings = spark_merged_clean.withColumn('mini-lm-vectors', embedding_udf(spark_merged_clean['Name']))\n",
     "\n",
-    "spark_embedded.select(\"Name\", \"NameEmbedding\").show(truncate=False)\n"
+    "# Show the result\n",
+    "df_with_embeddings.head()\n",
+    "\n"
    ]
   }
  ],