diff --git a/.github/workflows/notebooks.yaml b/.github/workflows/notebooks.yaml
index f42843f9b..3c7f96e90 100644
--- a/.github/workflows/notebooks.yaml
+++ b/.github/workflows/notebooks.yaml
@@ -1,5 +1,3 @@
-# Notebook-related checks
-
 name: Notebooks
 
 on:
@@ -17,11 +15,49 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
+      with:
+        # Checkout PR branch for PR events, main branch for manual runs
+        ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || 'main' }}
+        fetch-depth: 0  # Fetch all history to avoid issues with rebasing
     - uses: actions/setup-python@v4
     - name: Install tensorflow-docs
       run: python3 -m pip install -U git+https://github.com/tensorflow/docs
     - name: Fetch main branch
       run: git fetch -u origin main:main
+    - name: Get changed notebooks
+      id: changed-notebooks
+      run: |
+        if [ "${{ github.event_name }}" == "pull_request" ]; then
+          # Only check notebooks modified in this pull request
+          git diff --name-only origin/main | grep '\.ipynb$' > changed_notebooks.txt
+        else
+          # Manual run, check everything
+          find . -name '*.ipynb' > changed_notebooks.txt
+        fi
+        echo "changed_notebooks=$(cat changed_notebooks.txt)" >> $GITHUB_OUTPUT
+    - name: Format notebooks
+      run: |
+        if [ -s "${{ steps.changed-notebooks.outputs.changed_notebooks }}" ]; then
+          echo "Formatting notebooks with nbfmt:"
+          python3 -m tensorflow_docs.tools.nbfmt ${{ steps.changed-notebooks.outputs.changed_notebooks }}
+        else
+          echo "No notebooks to format."
+        fi
+    - name: Commit and push changes (if any)
+      if: ${{ success() }}
+      run: |
+        git config --local user.email "action@github.com"
+        git config --local user.name "GitHub Action"
+        # Check if there are changes
+        if ! git diff --quiet; then
+          # Commit on the current branch
+          git add .
+          git commit -m "Format notebooks"
+          # Push directly to the PR branch
+          git push origin ${{ github.event.pull_request.head.ref }}
+        else
+          echo "No changes to commit."
+        fi
     - name: Check notebook formatting
       run: |
         if [ "${{ github.event_name }}" == "pull_request" ]; then
@@ -38,7 +74,6 @@ jobs:
           echo "Check formatting with nbfmt:"
           python3 -m tensorflow_docs.tools.nbfmt --test "${changed_notebooks[@]}"
         fi
-
   nblint:
     name: Notebook lint
     runs-on: ubuntu-latest
@@ -75,5 +110,4 @@ jobs:
             --arg=base_url:https://ai.google.dev/ \
             --exclude_lint=tensorflow::button_github \
             "${changed_notebooks[@]}"
-        fi
-
+        fi
\ No newline at end of file
diff --git a/gemini-2/websockets/live_api_starter.ipynb b/gemini-2/websockets/live_api_starter.ipynb
index b75c75c73..d864e7d79 100644
--- a/gemini-2/websockets/live_api_starter.ipynb
+++ b/gemini-2/websockets/live_api_starter.ipynb
@@ -74,9 +74,7 @@
     },
     {
       "cell_type": "markdown",
-      "metadata": {
-        "id": "gnzxta_57_Ip"
-      },
+      "metadata": {},
       "source": [
         "## Setup"
       ]
@@ -451,9 +449,6 @@
       },
       "outputs": [
         {
-          "metadata": {
-            "tags": null
-          },
           "name": "stdout",
           "output_type": "stream",
           "text": [
@@ -479,9 +474,6 @@
           "output_type": "display_data"
         },
         {
-          "metadata": {
-            "tags": null
-          },
           "name": "stdout",
           "output_type": "stream",
           "text": [
@@ -506,9 +498,6 @@
           "output_type": "display_data"
         },
         {
-          "metadata": {
-            "tags": null
-          },
           "name": "stdout",
           "output_type": "stream",
           "text": [
@@ -533,9 +522,6 @@
           "output_type": "display_data"
         },
         {
-          "metadata": {
-            "tags": null
-          },
           "name": "stdout",
           "output_type": "stream",
           "text": [
@@ -560,9 +546,6 @@
           "output_type": "display_data"
         },
         {
-          "metadata": {
-            "tags": null
-          },
           "name": "stdout",
           "output_type": "stream",
           "text": [