{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Copyright 2026 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "# Generate descriptions from videos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment Setup (Antigravity / Data Cloud)\n",
    "\n",
    "If you are running this notebook in **Antigravity** (with the Google Data Cloud extension), the Spark session (`spark` variable) is automatically initialized once you select a remote kernel. \n",
    "\n",
    "If you are running **locally** or in another environment, you will need to manually initialize the `SparkSession` and install any missing dependencies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Overview\n",
    "\n",
    "This notebook shows how to generate descriptions of videos in a GCS bucket.  \n",
    "It uses the [Youtube UGC dataset](https://media.withyoutube.com/) and uses the [Gemini](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/gemini) to generate video descriptions for each video."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### **Steps**\n",
    "Using Spark,\n",
    "1) It reads the table [Youtube UGC dataset](https://media.withyoutube.com/) from gs://dataproc-metastore-public-binaries/youtube_ucg/\n",
    "2) It calls Vertex AI Gemini API vision pro to generate description from videos."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from pyspark.sql.functions import regexp_replace, concat\n",
    "from pyspark.sql.functions import udf, col, lit\n",
    "\n",
    "import google.auth\n",
    "import google.auth.transport.requests\n",
    "import requests\n",
    "\n",
    "import pandas as pd\n",
    "pd.set_option('display.max_colwidth', None)\n",
    "pd.set_option('display.min_rows', 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Read dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "BINARIES_BUCKET_PATH = \"gs://dataproc-metastore-public-binaries/youtube_ucg/\"\n",
    "binaries_df = spark.read.format(\"binaryFile\").option(\"recursiveFileLookup\", \"true\").load(BINARIES_BUCKET_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Let's select the paths of the first 5 youtube videos\n",
    "paths_df = binaries_df.select(\"path\").limit(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Define UDF and call Gemini API to generate video descriptions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_schema = {\n",
    "    \"type\": \"object\",\n",
    "    \"properties\": {\n",
    "        \"where\": {\"type\": \"string\"},\n",
    "        \"how_many_people\": {\"type\": \"integer\"},\n",
    "        \"task\": {\"type\": \"string\"},\n",
    "        \"proposition\": {\"type\": \"string\"},\n",
    "        \"description\": {\"type\": \"string\"}\n",
    "    },\n",
    "    \"required\": [\"where\",\"how_many_people\",\"task\",\"proposition\",\"description\"],\n",
    "}\n",
    "\n",
    "system_instructions = [\n",
    "    \"\"\"Format the 5 items as attributes of a JSON object: where, how_many_people, task, proposition and description.\"\"\",  \n",
    "    \"\"\"The response should be a single valid formatted JSON object only.\"\"\"\n",
    "]\n",
    "\n",
    "prompt = f\"\"\"\n",
    "    Create a short description for this video with the following questions:\n",
    "     1) Where is the video recorded? \n",
    "     2) How many people are shown in the video? \n",
    "     3) What the people are doing in the video? \n",
    "     4) Whats the proposition for the video, i.e what it is about?\n",
    "     5) A sumary description from the itens 1,2,3 and 4\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vertexai.generative_models import GenerativeModel, GenerationConfig, Part, Image, Content, HarmCategory, HarmBlockThreshold\n",
    "\n",
    "def predict(uri, prompt, system_instructions=system_instructions, response_schema=response_schema, content_type=\"video/mp4\", temperature=1, model_name=\"gemini-2.5-pro\"):\n",
    "\n",
    "    model = GenerativeModel(model_name=model_name, system_instruction=system_instructions)\n",
    "    \n",
    "    prompt_content = Content(\n",
    "        role=\"user\",\n",
    "        parts=[\n",
    "            Part.from_uri(uri, content_type),\n",
    "            Part.from_text(prompt)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    response = model.generate_content(\n",
    "        prompt_content,\n",
    "        generation_config = GenerationConfig(\n",
    "            temperature=temperature, response_mime_type=\"application/json\", response_schema=response_schema\n",
    "        ),\n",
    "        safety_settings={\n",
    "                HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_ONLY_HIGH\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    return response.text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def generate_descriptions(gcs_uri):\n",
    "\n",
    "    descriptions = predict(gcs_uri, prompt)\n",
    "    return descriptions\n",
    "    \n",
    "generate_descriptions_udf = udf(generate_descriptions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "df_descriptions = paths_df.sort(paths_df.path.asc()).withColumn(\"data\", generate_descriptions_udf(paths_df.path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_descriptions.cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_descriptions.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extract feature from generated text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import from_json, col\n",
    "from pyspark.sql.types import StructType, StructField, StringType\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField('where', StringType(), True),\n",
    "        StructField('how_many_people', StringType(), True),\n",
    "        StructField('proposition', StringType(), True),\n",
    "        StructField('description', StringType(), True),\n",
    "        StructField('task', StringType(), True)\n",
    "    ]\n",
    ")\n",
    "df_final = df_descriptions.withColumn(\"exploded_data\", from_json(regexp_replace(regexp_replace(col(\"data\"),\"json\", \"\"),\"```\",\"\"), schema))\\\n",
    "    .select(col('path'),col('exploded_data.*'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "df_final.toPandas()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "runtime-iceberg on Serverless Spark",
   "language": "python",
   "name": "9c39b79e5d2e7072beb4bd59-runtime-iceberg"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}